Error when learning rate is big

[f-cg]

python3 main.py --max-depth 4 --lr 1
...
Train Epoch: 3 [17280/60000 (29%)] Loss: 0.522095, Accuracy: 54/64 (84.0000%)
Traceback (most recent call last):
File "main.py", line 79, in
model.train_(train_loader, epoch)
File "/content/model.py", line 190, in train_
correct += pred.eq(target.data).cpu().sum()
RuntimeError: Expected object of backend CPU but got backend CUDA for argument #2 'other'

This is so stange and I can't fix it.

[kimhc6028]

do you need high lr? does this bug occur when lr is relatively small? lr=1 seems like unlikely

[f-cg]

Ok, I don't need that high lr, but this issue can happen when lr is smaller and epoch is bigger.