| 1 | import torch |
| 2 | |
| 3 | def optimizer_to(optim, device): |
| 4 | for param in optim.state.values(): |
| 5 | # Not sure there are any global tensors in the state dict |
| 6 | if isinstance(param, torch.Tensor): |
| 7 | param.data = param.data.to(device) |
| 8 | if param._grad is not None: |
| 9 | param._grad.data = param._grad.data.to(device) |
| 10 | elif isinstance(param, dict): |
| 11 | for subparam in param.values(): |
| 12 | if isinstance(subparam, torch.Tensor): |
| 13 | subparam.data = subparam.data.to(device) |
| 14 | if subparam._grad is not None: |
| 15 | subparam._grad.data = subparam._grad.data.to(device) |