package formatting for easier use in the future ?

windsornguyen / flash-stu

PyTorch implementation of the Spectral Transform Unit.

https://arxiv.org/abs/2409.10489/

Apache License 2.0

9 stars 2 forks source link

package formatting for easier use in the future ? #2

Closed hubstrauss closed 1 month ago

hubstrauss commented 1 month ago

As i wanted to re-use this Flash STU within my code, I did a quick refactoring so that it's easier for people to import and use the module. Haven't tested the training though - but I thought it could be a useful improvement for everyone Thanks for the nice paper and repo btw !