Hi, I am trying to reproduce the [notebook](https://github.com/NVIDIA/torch-harmonics/blob/main/notebooks/train_sfno.ipynb) example but with galewsky initial condition. The losses during training have nan in them Can you help me with this error? P.S. : I can attach the notebook if you want it.
