code: def move_optimizer_to_cuda(optimizer):
      
      for param_group in optimizer.param_groups:
          for param in param_group['params']:
              if param.is_cuda:
                  param_state = optimizer.state[param]
                  for k in param_state.keys():
                      if isinstance(param_state[k], torch.Tensor):
                          param_state[k] = param_state[k].cuda(device=param.get_device())
docs: Move the optimizer state to GPU, if necessary.
      After calling, any parameter specific state in the optimizer
      will be located on the same device as the parameter.