Repository: pytorch/pytorch

Issue:
optimizer load_state_dict() problem?
Hi, I encountered this bug:
```

    optimizer.step()
    exp_avg.mul_(beta1).add_(1 - beta1, grad)

TypeError: add_ received an invalid combination of arguments - got (float, torch.cuda.FloatTensor), but expected one of:
 * (float value)
 * (torch.FloatTensor other)
 * (torch.SparseFloatTensor other)
 * (float value, torch.FloatTensor other)
      didn't match because some of the arguments have invalid types: (float, torch.cuda.FloatTensor)
 * (float value, torch.SparseFloatTensor other)
      didn't match because some of the arguments have invalid types: (float, torch.cuda.FloatTensor)
```

The code skeleton is like:
```
model = Model()
model.load_state_dict(checkpoint['model'])
model.cuda()

optimizer = optim.Adam()
optimizer.load_state_dict(checkpoint['optimizer'])

...
#  In train loop
for epoch in range(...):
  ...
  optimizer.step()
     -> BUG <-
```

It seems the loaded  `param_groups` are `torch.cuda.FloatTensor`, and I've tried a workaround to
move `optmizer.param_groups` to `cpu`, but it still has the same bug.


Discussion:
Could you provide a full script to reproduce the problem?
maybe you can try  like this,
      optimizer.step()
      exp_avg.mul_(beta1).add_(1 - beta1, grad.cpu())

Sorry,  I missed the reply email.

I am afraid that I am unable to provide a reproducer now.  It is a work I am doing for the OpenNMT-py project:https://github.com/OpenNMT/OpenNMT-py, trying to use `lr_scheduler` for doing lr update. And I encoutered this problem when testing the `resume a suspended training` case.  So I factored out the code skeleton about this problem above.

I've tried several methods, including tricks like what @hefeicyp suggests, but it still happens.

Per my analysis, it is because the previous training was done on gpu, so when saving the `optimizer.state_dict`, the stored states(tensors) are of `cuda` version. During resuming, when we load the saved  optimizer, `load_state_dict()` loads this `cuda` version to cpu(the model(`nn.Module`) can be moved to gpu easily, but `torch.optimizer` seems lacking this ability?) , so this problem emerges.   
Try moving optimizer state to the GPU memory manually after loading it from the checkpoint.
```python
optimizer = optim.Adam()
optimizer.load_state_dict(checkpoint['optimizer'])
for state in optimizer.state.values():
    for k, v in state.items():
        if isinstance(v, torch.Tensor):
            state[k] = v.cuda()
```
I agree that having an `optimizer.cuda()` method for this operation would be nice.
@dogancan, thanks.  My work was suspended due to other problems, when resumed, I will try your method. 
I'm afraid @dogancan's solution won't work. It will make the error go away, but your optimizer will no longer be training the model. You should recreate optimizers after casting modules to a different type or device, and you can use `load_state_dict` to restore the state from a previous copy. This currently doesn't work, but we should fix it (by copying from the data from the state dict, instead of using the tensors directly - this allows for cross-device or cross-type updates).
@apaszke , yep,  your method is what I currently use, it works.  But I will wait for upstream to fix this problem though.  Thanks for your great works!
@apaszke Ah, my bad. I forgot to update the line where the optimizer is recreated. But otherwise, the following should do the job, right?
```python
model = Model()
model.load_state_dict(checkpoint['model'])
model.cuda()
optimizer = optim.Adam(model.parameters())
optimizer.load_state_dict(checkpoint['optimizer'])
for state in optimizer.state.values():
    for k, v in state.items():
        if isinstance(v, torch.Tensor):
            state[k] = v.cuda()
```

Resolving patch:
diff --git a/torch/optim/lbfgs.py b/torch/optim/lbfgs.py
--- a/torch/optim/lbfgs.py
+++ b/torch/optim/lbfgs.py
@@ -93,7 +93,9 @@ def step(self, closure):
         line_search_fn = group['line_search_fn']
         history_size = group['history_size']
 
-        state = self.state['global_state']
+        # NOTE: LBFGS has only global state, but we register it as state for
+        # the first param, because this helps with casting in load_state_dict
+        state = self.state[self._params[0]]
         state.setdefault('func_evals', 0)
         state.setdefault('n_iter', 0)
 
diff --git a/torch/optim/optimizer.py b/torch/optim/optimizer.py
--- a/torch/optim/optimizer.py
+++ b/torch/optim/optimizer.py
@@ -1,4 +1,4 @@
-from collections import defaultdict
+from collections import defaultdict, Iterable
 
 import torch
 from copy import deepcopy
@@ -96,8 +96,33 @@ def load_state_dict(self, state_dict):
         id_map = {old_id: p for old_id, p in
                   zip(chain(*(g['params'] for g in saved_groups)),
                       chain(*(g['params'] for g in groups)))}
-        state = defaultdict(
-            dict, {id_map.get(k, k): v for k, v in state_dict['state'].items()})
+
+        def cast(param, value):
+            """Make a deep copy of value, casting all tensors to device of param."""
+            if torch.is_tensor(value):
+                # Floating-point types are a bit special here. They are the only ones
+                # that are assumed to always match the type of params.
+                if any(tp in type(param.data).__name__ for tp in {'Half', 'Float', 'Double'}):
+                    value = value.type_as(param.data)
+                value = value.cuda(param.get_device()) if param.is_cuda else value.cpu()
+                return value
+            elif isinstance(value, dict):
+                return {k: cast(param, v) for k, v in value.items()}
+            elif isinstance(value, Iterable):
+                return type(value)(cast(param, v) for v in value)
+            else:
+                return value
+
+        # Copy state assigned to params (and cast tensors to appropriate types).
+        # State that is not assigned to params is copied as is (needed for
+        # backward compatibility).
+        state = defaultdict(dict)
+        for k, v in state_dict['state'].items():
+            if k in id_map:
+                param = id_map[k]
+                state[param] = cast(param, v)
+            else:
+                state[k] = v
 
         # Update parameter groups, setting their 'params' value
         def update_group(group, new_group):


Test patch:
diff --git a/test/test_optim.py b/test/test_optim.py
--- a/test/test_optim.py
+++ b/test/test_optim.py
@@ -133,7 +133,8 @@ def _test_state_dict(self, weight, bias, input, constructor):
 
         def fn_base(optimizer, weight, bias):
             optimizer.zero_grad()
-            loss = (weight.mv(input) + bias).pow(2).sum()
+            i = input_cuda if weight.is_cuda else input
+            loss = (weight.mv(i) + bias).pow(2).sum()
             loss.backward()
             return loss
 
@@ -161,6 +162,29 @@ def fn_base(optimizer, weight, bias):
         # Make sure state dict wasn't modified
         self.assertEqual(state_dict, state_dict_c)
 
+        # Check that state dict can be loaded even when we cast parameters
+        # to a different type and move to a different device.
+        if not torch.cuda.is_available():
+            return
+
+        input_cuda = Variable(input.data.float().cuda())
+        weight_cuda = Variable(weight.data.float().cuda(), requires_grad=True)
+        bias_cuda = Variable(bias.data.float().cuda(), requires_grad=True)
+        optimizer_cuda = constructor(weight_cuda, bias_cuda)
+        fn_cuda = functools.partial(fn_base, optimizer_cuda, weight_cuda, bias_cuda)
+
+        state_dict = deepcopy(optimizer.state_dict())
+        state_dict_c = deepcopy(optimizer.state_dict())
+        optimizer_cuda.load_state_dict(state_dict_c)
+        # Make sure state dict wasn't modified
+        self.assertEqual(state_dict, state_dict_c)
+
+        for i in range(20):
+            optimizer.step(fn)
+            optimizer_cuda.step(fn_cuda)
+            self.assertEqual(weight, weight_cuda)
+            self.assertEqual(bias, bias_cuda)
+
     def _test_basic_cases(self, constructor, ignore_multidevice=False):
         self._test_state_dict(
             torch.randn(10, 5),