Unable to load VGG model's state dict on GPU

[harshbafna]

I am trying a generic code to load the state-dict of a model on CPU/GPU. It works fine for other vision type models but fails for VGG models on GPU machines

Here is a sample code I am trying out :

import torch
map_location = 'cuda' if torch.cuda.is_available() else 'cpu'
model_pt_path = 'vgg13-c768596a.pth'
state_dict = torch.load(model_pt_path, map_location=map_location)

The above code fails with following error

Traceback (most recent call last):
  File "test_vgg.py", line 4, in <module>
    state_dict = torch.load(model_pt_path, map_location=map_location)
  File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/torch/serialization.py", line 593, in load
    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)
  File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/torch/serialization.py", line 773, in _legacy_load
    result = unpickler.load()
  File "/home/ubuntu/anaconda3/lib/python3.7/site-packages/torch/tensor.py", line 147, in __setstate__
    self.set_(*state)
RuntimeError: Expected object of device type cpu but got device type cuda for argument #2 'source'

I tried this same code with Densenet/Alexnet/Squeezenet models and it works fine on both GPU and CPU machines

Environment information :

OS : Ubuntu 18.0.4
PyTorch Version : 1.5.1
TorchVision Version : 0.6.1

[pmeier]

To narrow this down a little: the error only happens for the vgg11 and vgg13 weights.

[fmassa]

This relates to #2068 although the error message is different, but the root issue is the same (very old serialization format).

The workaround on your side for now could be to avoid using the map_location for CUDA, and convert manually the tensors to GPU if needed (or maybe even just pass the weights on CPU to load_state_dict while the model is on CUDA, I think this will also work).

[NicolasHug]

This now fails with

Traceback (most recent call last):
  File "/fsx/users/nicolashug/work/vision/lol.py", line 21, in <module>
    state_dict = torch.load(model_pt_path, map_location=map_location)
  File "/fsx/users/nicolashug/conda/envs/pt/lib/python3.9/site-packages/torch/serialization.py", line 607, in load
    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)
  File "/fsx/users/nicolashug/conda/envs/pt/lib/python3.9/site-packages/torch/serialization.py", line 786, in _legacy_load
    result = unpickler.load()
  File "/fsx/users/nicolashug/conda/envs/pt/lib/python3.9/site-packages/torch/_tensor.py", line 189, in __setstate__
    self.set_(*state)
RuntimeError: Attempted to set the storage of a tensor on device "cpu" to a storage on different device "cuda:0".  This is no longer allowed; the devices must match.

which is similar to #3767

I'll try to update the weights of VGG as done for the other models in #3620, since it seems to have solved the other issue #3767 (comment) (EDIT: it also seems to solve this specific issue for VGG)

[NicolasHug]

BTW I can only reproduce the issue for inception_v3_google-9aff3935.pth, vgg11-7de3a66a.pth, vgg13-b20e5c87.pth
The rest of the VGG weights seem fine.