From 6ca0f1f2fd02a6d9405c1fd99fdd6c504d0a9038 Mon Sep 17 00:00:00 2001 From: Lu Fang Date: Thu, 7 Sep 2017 15:01:34 -0700 Subject: [PATCH] Load the weights to CPU if CUDA is not available --- advanced_source/super_resolution_with_caffe2.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/advanced_source/super_resolution_with_caffe2.py b/advanced_source/super_resolution_with_caffe2.py index a86f9c5a543..4ba84299eab 100644 --- a/advanced_source/super_resolution_with_caffe2.py +++ b/advanced_source/super_resolution_with_caffe2.py @@ -89,7 +89,10 @@ def _initialize_weights(self): batch_size = 1 # just a random number # Initialize model with the pretrained weights -torch_model.load_state_dict(model_zoo.load_url(model_url)) +map_location = lambda storage, loc: storage +if torch.cuda.is_available(): + map_location = None +torch_model.load_state_dict(model_zoo.load_url(model_url, map_location=map_location)) # set the train mode to false since we will only run the forward pass. torch_model.train(False)