Compatibility issues for pytorch_lightning

[qiyang77]

Taking the N-body minist script as an example, the code can't be executed with pytorch 2.0.1 due to the compatibility issue of pytorch_lightning.

Is there any workarounds to bypass these issues without installing an old version torch? Thx!

For pytorch_lightning=2.2.1

  File ~\Anaconda3\envs\neuraloperator\lib\site-packages\pytorch_lightning\utilities\argparse.py:70 in insert_env_defaults
    return fn(self, **kwargs)

TypeError: Trainer.__init__() got an unexpected keyword argument 'track_grad_norm'

For pytorch_lightning=1.9.0

File ~\Anaconda3\envs\neuraloperator\lib\site-packages\pytorch_lightning\utilities\rank_zero.py:69 in rank_zero_debug

NameError: name '_debug' is not defined

For pytorch_lightning=1.6.4

MisconfigurationException: The provided lr scheduler `SequentialLR` doesn't follow PyTorch's LRScheduler API. You should override the `LightningModule.lr_scheduler_step` hook with your own logic if you are using a custom LR scheduler.

[gaozhihan]

Thank you for reporting this issue. The compatibility between different versions of PyTorch Lightning (PL) can be troublesome, even between minor version changes.
There is a straightforward solution that avoids such dependency issues. Notice that the Earthformer model and the pretrained weights are purely in PyTorch, without the dependencies on PL. You could implement your own training script with your specific version of PL.
You may want to refer to the following scripts for guidance on how to use Earthformer purely in PyTorch:

earth-forecasting-transformer/tests/unittests/test_pretrained_checkpoints.py

Line 28 in 7732b03

def config_cuboid_transformer(cfg, model_type="CuboidTransformerModel"):

and

https://github.com/gaozhihan/earth-forecasting-transformer/blob/a5c07f22ec53ba577d679e0a3be8eb7e77d3e82c/tests/test_cuboid.py