Training with Lightning>=2.0.0 and Ubuntu 22.04

[Achleshwar]

Hi @nihalsid, thank you for this amazing work.

Issue

1. Ubuntu 22.04 might face issues with the current environment setup (requirements.txt) - as cudatoolkit 11.3 is not available for ubuntu 22.04.
2. I also wanted to modify the current codebase to make it work with the new lightning versions since there have been significant modifications after the release of Lightning 2.0.0

Environment changes

conda create -n panopli python=3.10

# check compatibility between torch lightning and pytorch versions
# the latest lightning version 2.2 is compatible with the following 
# torch: >=1.13; <=2.2
# torchmetrics: >= 0.7.0
# python: >=3.8; <=3.11

pip install torch==2.2.0 torchvision==0.17.0 torchaudio==2.2.0 --index-url https://download.pytorch.org/whl/cu118
pip install torchmetrics
python -m pip install lightning
pip install torch-scatter -f https://data.pyg.org/whl/torch-2.2.0+cu118.html

# comment the above downloaded packages in requirements.txt
pip install -r requirements.txt

Code modifications

There are a couple of changes required in trainer directory. Please feel free to check the modified file at this gist for train_panopli_tensorf.py and this gist for __init__.py.

Major changes in trainer/train_panopli_tensorf.py are listed below

• validation_step

• validation_epoch_end becomes on_validation_epoch_end

• 'on_validation_epoch_end`