Suggestion to handle model deserialization on CPU-only machines

[thekingofall]

I recently encountered an issue when using your software on a CPU-only machine. The problem occurs when attempting to load a model that was trained on a CUDA device while running the code on a machine without GPU support. The specific error message is:

warnings.warn(w)
Traceback (most recent call last):
File "Predict.py", line 103, in
df = predict(sys.argv[1], sys.argv[2])
File "Predict.py", line 85, in predict
model, train_file = get_model(dataset)
File "Predict.py", line 71, in get_model
model = load_model(hparams, checkpoint)
File "Predict.py", line 42, in load_model
model = ERGOLightning(hparams)
...
To make your software more versatile and compatible with CPU-only machines, I suggest modifying the line in Models.py where the model is being loaded:

checkpoint = torch.load(ae_file)

Change it to:

checkpoint = torch.load(ae_file, map_location=torch.device('cpu'))

By adding the map_location parameter, the model will be loaded on the CPU even if it was originally trained on a CUDA device. This small change will make your software more accessible to users without GPU support.

Thank you for your attention to this matter, and I hope this suggestion is helpful.