_optimizer_to_ now correctly send to the input device

Author: Aethor

tibert/train.py

@@ -75,7 +75,7 @@ def _optimizer_to_(
     for state in optimizer.state.values():
         for k, v in state.items():
             if isinstance(v, torch.Tensor):
-                state[k] = v.cuda()
+                state[k] = v.to(device)
     return optimizer