Explain how to train on GPU and load on CPU

[hrzn]

[dennisbader]

Here is a code snippet (I couldn't test on GPU but I believe it should work):

# create, train, and save a model on GPU (use a single GPU, the first one with "gpus"=[0])
gpu_params = {
    "accelerator": "gpu",
    "gpus": [0],  # use "devices" instead of "gpus" for PyTorch Lightning >= 1.7.
    "auto_select_gpu": False,
}
model = NBEATSModel(..., pl_trainer_kwargs=gpu_params)
model.fit(...)
model.save("model.pt")

# to load a model on CPU that was saved on GPU, we simply have to update the trainer parameters
# to map to the CPU device
model_loaded = NBEATSModel.load("model.pt")
cpu_params = {
    "accelerator": "cpu",
    "gpus": None,  # use "devices" instead of "gpus" for PyTorch Lightning >= 1.7.
    "auto_select_gpus": False
}
model_loaded.trainer_params.update(cpu_params)
preds = model.predict(...)

[hrzn]

@dennisbader could we update this example with the latest version of save/load functions?