MPS with num_workers fast but checkpoints don't work

[velaia]

I've been training a ResNet50 with custom head on a M1 Mac using a Jupyter notebook in Jupyter Lab. To accelerate the training I tried the num_workers argument in trainer.fit(). It did not work in Jupyter but when I created a Python file and put the code in the main() method, the training ran with a speedup of about 4x.

Unfortunately when I load the model from a checkpoint the loss is like I haven't trained the model at all! This only happens when I add num_workers.

I'm not sure it's a Lightning bug or issue or related to PyTorch, that's why I put it here in the discussion section instead of filing an issue.

My plan is to discard MPS for training as I lose a lot of time with MPS-related issues in many projects. Posting it here for others that might run into the issue mostly.