Error deploying pretrained PyTorch model to sagemaker

Running into issue loading the model. I have defined my model and am trying to deploy it on a non GPU instance.
Not sure where the error is since i’m specifying it to load it on CPU. Should this not be the case ? I have tried passing on torch.device to map to location as well.

The model loads just fine when i’m working remotely, but fails during the deployment.
Here are the error Logs and my model_fn used.

PyTorch version: 1.1.0

Error Logs

Traceback (most recent call last):
  File "/usr/local/lib/python3.6/dist-packages/gunicorn/workers/", line 56, in handle
    self.handle_request(listener_name, req, client, addr)
  File "/usr/local/lib/python3.6/dist-packages/gunicorn/workers/", line 160, in handle_request
  File "/usr/local/lib/python3.6/dist-packages/gunicorn/workers/", line 107, in handle_request
    respiter = self.wsgi(environ, resp.start_response)
  File "/usr/local/lib/python3.6/dist-packages/sagemaker_pytorch_container/", line 107, in main
  File "/usr/local/lib/python3.6/dist-packages/sagemaker_containers/", line 157, in initialize
    self._model = self._model_fn(_env.model_dir)
  File "/usr/local/lib/python3.6/dist-packages/sagemaker_containers/", line 87, in wrapper
    six.reraise(error_class, error_class(e), sys.exc_info()[2])
  File "/usr/local/lib/python3.6/dist-packages/", line 692, in reraise
    raise value.with_traceback(tb)
  File "/usr/local/lib/python3.6/dist-packages/sagemaker_containers/", line 85, in wrapper
    return fn(*args, **kwargs)
  File "/opt/ml/code/", line 128, in model_fn
    model.load_state_dict(torch.load(f), map_to_location='cpu')
  File "/usr/local/lib/python3.6/dist-packages/torch/", line 387, in load
    return _load(f, map_location, pickle_module, **pickle_load_args)
  File "/usr/local/lib/python3.6/dist-packages/torch/", line 574, in _load
    result = unpickler.load()
  File "/usr/local/lib/python3.6/dist-packages/torch/", line 537, in persistent_load
    deserialized_objects[root_key] = restore_location(obj, location)
  File "/usr/local/lib/python3.6/dist-packages/torch/", line 119, in default_restore_location
    result = fn(storage, location)
  File "/usr/local/lib/python3.6/dist-packages/torch/", line 95, in _cuda_deserialize
    device = validate_cuda_device(location)
  File "/usr/local/lib/python3.6/dist-packages/torch/", line 79, in validate_cuda_device
    raise RuntimeError('Attempting to deserialize object on a CUDA '


def model_fn(model_dir):
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')'Loading the model.')
    model = Generator()
    if torch.cuda.is_available():
        with open(os.path.join(model_dir, 'model.pth'), 'rb') as f:
        with open(os.path.join(model_dir, 'model.pth'), 'rb') as f:
            model.load_state_dict(torch.load(f), map_to_location='cpu')'Done loading model')
    return model

Deployment Script

role = get_execution_role() 
# You can also configure a sagemaker role and reference it by its name. 
# role = "CustomSageMakerRoleName"

pytorch_model = PyTorchModel(
predictor = pytorch_model.deploy(instance_type='ml.t2.medium', initial_instance_count=1,)

Source: Python Questions