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

Viewed 1231

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!

3 Answers

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

Can you tell me which Optimizer you are using. I have encountered this with AdamW optimizer. You can avoid it by loading optimizer with the load_state_dict and then mapping it to cpu explicitly using .cpu() function.

If you are using pytorch 1.12.0 with cuda binaries 11.6/11.7 then on your shell or command prompt, paste the following;

pip install torch==1.12.1+cu116 torchvision==0.13.1+cu116 torchaudio==0.12.1 --extra-index-url https://download.pytorch.org/whl/cu116

The Adam Optimizer regression was removed in the updated torch version

Related