Collectives™ on Stack Overflow

Find centralized, trusted content and collaborate around the technologies you use most.

Teams

Q&A for work

Connect and share knowledge within a single location that is structured and easy to search.

I am using pytorch to train a Neural Network. When I train and test on GPU, it works fine. But When I try to load the model parameters on CPU using:

net.load_state_dict(torch.load('rnn_x_epoch.net'))
I get the following error:
RuntimeError: cuda runtime error (35) : CUDA driver version is insufficient for CUDA runtime version at torch/csrc/cuda/Module.cpp:51
I have searched for the error, it's mainly because of CUDA driver dependency, but since I'm running on CPU when I get this error,it must be something else, or may be I missed something.
Since it's working fine using GPU, I could just run it on GPU but I'm trying to train the network on GPU, store the parameters and then load it on CPU mode for predictions.
I am just looking for a way to load the parameters while on CPU mode.
I tried this as well to load the parameters:
check = torch.load('rnn_x_epoch.net')
It did not work.
I tried to save the model parameters in two ways, to see if any of these would work, but didn't:
checkpoint = {'n_hidden': net.n_hidden,
          'n_layers': net.n_layers,
          'state_dict': net.state_dict(),
          'tokens': net.chars}
with open('rnn_x_epoch.net', 'wb') as f:
    torch.save(checkpoint, f)
torch.save(model.state_dict(), 'rnn_x_epoch.net')
TraceBack:
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-9-e61f28013b35> in <module>()
----> 1 net.load_state_dict(torch.load('rnn_x_epoch.net'))
/opt/conda/lib/python3.6/site-packages/torch/serialization.py in load(f, map_location, pickle_module)
    301         f = open(f, 'rb')
    302     try:
--> 303         return _load(f, map_location, pickle_module)
    304     finally:
    305         if new_fd:
/opt/conda/lib/python3.6/site-packages/torch/serialization.py in _load(f, map_location, pickle_module)
    467     unpickler = pickle_module.Unpickler(f)
    468     unpickler.persistent_load = persistent_load
--> 469     result = unpickler.load()
    471     deserialized_storage_keys = pickle_module.load(f)
/opt/conda/lib/python3.6/site-packages/torch/serialization.py in persistent_load(saved_id)
    435             if root_key not in deserialized_objects:
    436                 deserialized_objects[root_key] = restore_location(
--> 437                     data_type(size), location)
    438             storage = deserialized_objects[root_key]
    439             if view_metadata is not None:
/opt/conda/lib/python3.6/site-packages/torch/serialization.py in default_restore_location(storage, location)
     86 def default_restore_location(storage, location):
     87     for _, _, fn in _package_registry:
---> 88         result = fn(storage, location)
     89         if result is not None:
     90             return result
/opt/conda/lib/python3.6/site-packages/torch/serialization.py in _cuda_deserialize(obj, location)
     68     if location.startswith('cuda'):
     69         device = max(int(location[5:]), 0)
---> 70         return obj.cuda(device)
/opt/conda/lib/python3.6/site-packages/torch/_utils.py in _cuda(self, device, non_blocking, **kwargs)
     66         if device is None:
     67             device = -1
---> 68     with torch.cuda.device(device):
     69         if self.is_sparse:
     70             new_type = getattr(torch.cuda.sparse, 
self.__class__.__name__)
/opt/conda/lib/python3.6/site-packages/torch/cuda/__init__.py in __enter__(self)
    223         if self.idx is -1:
    224             return
--> 225         self.prev_idx = torch._C._cuda_getDevice()
    226         if self.prev_idx != self.idx:
    227             torch._C._cuda_setDevice(self.idx)
RuntimeError: cuda runtime error (35) : CUDA driver version is insufficient for CUDA runtime version at torch/csrc/cuda/Module.cpp:51
Also may be the save/load operations in Pytorch are only for GPU mode, but I am not really convinced by that.
  
When you call torch.load() on a file which contains GPU tensors, those tensors will be loaded to GPU by default. 
To load the model on CPU which was saved on GPU, you need to pass map_location argument as cpu in load function as follows:
# Load all tensors onto the CPU
net.load_state_dict(torch.load('rnn_x_epoch.net', map_location=torch.device('cpu')))
In doing so, the storages underlying the tensors are dynamically remapped to the CPU device using the map_location argument. You can read more on the official PyTorch tutorials.
This can also be done as follows:
# Load all tensors onto the CPU, using a function
net.load_state_dict(torch.load('rnn_x_epoch.net', map_location=lambda storage, loc: storage))
                These give the error as      TypeError: 'torch.Device' object is not callable. I'll look through the pytorch tutorials though!
– naman1994
                Jan 30, 2019 at 8:22
        
Thanks for contributing an answer to Stack Overflow!
Please be sure to answer the question. Provide details and share your research!
But avoid …
Asking for help, clarification, or responding to other answers.
Making statements based on opinion; back them up with references or personal experience.
To learn more, see our tips on writing great answers.