Trying to execute run.py in train folder renders an error

[bhardwaj-garvit]

Hi,
Really great and helpful code!
I was trying to run train.py on the nimrod-uk-1km-test data and encountered the following error, it says "RuntimeError: Serialization of parametrized modules is only supported through state_dict()." I searched on torch's website and found an earlier commit, so I downgraded torch to v1.12.0 but this did not go away.
Torch Link: pytorch/pytorch#69413

Can you guys help in debugging this issue? I am planning to use this on another dataset

**To Reproduce** Steps to reproduce the behavior: 1. installing dependencies 2. execute train/run.py and the above error shows in the terminal

[jacobbieker]

Hi, are you using multiple GPUs? By default the run.py tries to use 6 GPUs, although it should be changed to 1. The spectrally normalized layers in PyTorch don't seem to work in multi-GPU setting as far as I have been able to get them. So if you do change it to 1 GPU, it should start training

[bhardwaj-garvit]

I was earlier using cpu's, to sort the issue started using 1 gpu, but the training fills virtual memory of upto 200 GB(my system's limit) and the dataloader worker is killed. Can you suggest a way to bypass this.

[Chevolier]

I met the same issue, the memory keeps increasing to 256GB in the data loading process until it got killed by the system, any solution to solve this?

[Chevolier]

I met the same issue, the memory keeps increasing to 256GB in the data loading process until it got killed by the system, any solution to solve this?

Updates: My problem is solved by setting streaming=True in TFDataset as follows for my own dataset, by doing this, data are not first loaded into memory.

class TFDataset(torch.utils.data.dataset.Dataset):
def init(self, data_path, split):
super().init()
# self.reader = load_dataset(
# "openclimatefix/nimrod-uk-1km", "sample", split=split, streaming=True
# )
self.reader = load_dataset(data_path, split=split, streaming=True)