AdamW step device error
I have a model:
```
class M(nn.Module):
  def __init__(self, ...):
    self.A = nn.Parameter(...)
    self.B = nn.Parameter(...)
    self.C = torch.einsum(..., self.A, self.B)

  def forward(self, D):
    return func(self.C.to(D.device), D)
```

However, I'm getting the following error when training this model on GPU:
```
File "file.py", line 71, in optimizer_step
optimizer.step()
File ".../lib/python3.7/site-packages/torch/optim/lr_scheduler.py", line 67, in wrapper
return wrapped(*args, **kwargs)
File ".../lib/python3.7/site-packages/transformers/optimization.py", line 155, in step
exp_avg.mul_(beta1).add_(grad, alpha=1.0 - beta1)
RuntimeError: expected device cuda:0 but got device cpu
```

This is happening in 2.11.0. Am I doing something incorrectly here?