Error loading a saved model to run inference (using ddp_notebook strategy) about lightning HOT 1 OPEN

@carlos-havier As mentioned in the PyTorch Lightning documentation, when using ddp_notebook, the downside is:

"GPU operations such as moving tensors to the GPU or calling torch.cuda functions before invoking Trainer.fit is not allowed."

This means that there can be no CUDA tensors before calling Trainer.fit. By default, when training, PyTorch Lightning saves the state_dict of the trainer as CUDA when using GPU. So when load from checkpoint, CUDA is initialised. You can verify this with a simple check:

print(torch.cuda.is_initialized())

This can be placed before and after calling:

pl_model = LT_timm_model.load_from_checkpoint(def_log_chkpt)

You'll observe that CUDA is initialized when calling load_from_checkpoint, and once CUDA is initialized here, it cannot be re-initialized in a different context as required by ddp_notebook.

The Fix:

Use map_location as CPU when calling load_from_checkpoint:

pl_model = LT_timm_model.load_from_checkpoint(def_log_chkpt, map_location=torch.device('cpu'))

For your reference, I have added a few lines to debug based on your notebook here.

from lightning.