AssertionError: If capturable=False, state_steps should not be CUDA tensors

Question

I get this error while loading model weights of a previous epoch on Google colab. I'm using PyTorch version 1.12.0. I can't downgrade to a lower version as there are external libraries that Im using that require Pytorch 1.12.0

Thanks!

Answer 1

It seems related to a newly introduced parameter (capturable) for the Adam and AdamW optimizers. Currently two workarounds:

1. forcing capturable = True after loading the checkpoint optim.param_groups[0]['capturable'] = True. This seems to slow down the model training by approx. 10% (YMMV depending on the setup).
2. Reverting PyTorch back to previous versions (could be 1.11.0).

Source: https://github.com/pytorch/pytorch/issues/80809#issuecomment-1173481031