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

[jinqiupeter]

I got this error when running python main.py data/obama/ --workspace trial_obama/ -O --iters 125000 --finetune_lips --patch_size 32:

loss=0.0060 (0.0036), lr=0.000087: :   0% 4/6364 [00:01<2:21:59,  1.34s/it]Traceback (most recent call last):
  File "main.py", line 248, in <module>
    trainer.train(train_loader, valid_loader, max_epochs)
  File "/workspace/ER-Nerf/nerf_triplane/utils.py", line 983, in train
    self.train_one_epoch(train_loader)
  File "/workspace/ER-Nerf/nerf_triplane/utils.py", line 1252, in train_one_epoch
    self.scaler.step(self.optimizer)
  File "/opt/conda/lib/python3.8/site-packages/torch/cuda/amp/grad_scaler.py", line 338, in step
    retval = self._maybe_opt_step(optimizer, optimizer_state, *args, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/torch/cuda/amp/grad_scaler.py", line 285, in _maybe_opt_step
    retval = optimizer.step(*args, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/torch/optim/lr_scheduler.py", line 65, in wrapper
    return wrapped(*args, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/torch/optim/optimizer.py", line 109, in wrapper
    return func(*args, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
    return func(*args, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/torch/optim/adamw.py", line 161, in step
    adamw(params_with_grad,
  File "/opt/conda/lib/python3.8/site-packages/torch/optim/adamw.py", line 218, in adamw
    func(params,
  File "/opt/conda/lib/python3.8/site-packages/torch/optim/adamw.py", line 259, in _single_tensor_adamw
    assert not step_t.is_cuda, "If capturable=False, state_steps should not be CUDA tensors."
AssertionError: If capturable=False, state_steps should not be CUDA tensors.
loss=0.0060 (0.0036), lr=0.000087: :   0% 5/6364 [00:01<31:14,  3.39it/s]

Can someone please help?

[1flyingpig]

I have same question, do you know how to deal with it

[jinqiupeter]

@1flyingpig I fixed it. Just comment out the two lines (in my case line 259) in /opt/conda/lib/python3.8/site-packages/torch/optim/adamw.py that asserts capturable

[pengpengzi]

@1flyingpig I fixed it. Just comment out the two lines (in my case line 259) in /opt/conda/lib/python3.8/site-packages/torch/optim/adamw.py that asserts capturable

I have also resolved it following your method. I want to ask if this will have any impact on subsequent training and results, please translate this sentence into English.

[Oyiyi]

Please let me know if any impact :) @pengpengzi @1flyingpig thanks!