Adding the non-spectral SFNO from nvidia #101
Conversation
|
|
||
| # classes | ||
|
|
||
| class CrossFormer(BaseModel): |
There was a problem hiding this comment.
Can you call the class CrossFormerGated or CrossFormerGRU?
There was a problem hiding this comment.
oops did not mean to add xformer-gated; was mainly testing Arnolds approach with the code base at the time. Will remove eventually
|
|
||
| # cube embedding | ||
|
|
||
| class CubeEmbedding(nn.Module): |
There was a problem hiding this comment.
For some of these layers that may be shared across multiple models, can we create a separate layers module or directory within credit/models?
| # pad_lat=pad_lat | ||
| # ).to("cuda") | ||
|
|
||
| image_height = 640 # 640, 192 |
There was a problem hiding this comment.
contents of if __name__ == "__main__ block should be moved to a unit test.
| import torch | ||
| import warnings | ||
| from tltorch.factorized_tensors.core import FactorizedTensor | ||
| import tensorly as tl |
There was a problem hiding this comment.
Tensorly and tltorch need to be added to pyproject.toml, requirements.txt, and the various environment.yml files.
There was a problem hiding this comment.
I am going to try to not depend on these packages; will update eventually
| return x | ||
|
|
||
|
|
||
| class SphericalFourierNeuralOperatorNet(nn.Module): |
There was a problem hiding this comment.
Docstring should be converted from numpy to google style. Also, you should add a link in the docstring to where the code originally came from.
|
|
||
|
|
||
| if __name__ == "__main__": | ||
| image_height = 192 # 640, 192 |
There was a problem hiding this comment.
Please create a unit test in the tests directory with this code.
Environment.yml has not bee updated yet