Allow using torch.device for loading

EthanSteinberg · EthanSteinberg · commit ec6af39eb1f2 · 2018-05-07T07:58:07.000-07:00
diff --git a/test/test_torch.py b/test/test_torch.py
@@ -6265,14 +6265,32 @@ def load_bytes():
             return data
 
         fileobject_lambdas = [lambda: test_file_path, load_bytes]
-        map_locations = [map_location, {'cuda:0': 'cpu'}, 'cpu']
+        map_locations = [map_location, {'cuda:0': 'cpu'}, 'cpu', torch.device('cpu')]
 
         for fileobject_lambda in fileobject_lambdas:
             for map_location in map_locations:
                 tensor = torch.load(fileobject_lambda(), map_location=map_location)
                 self.assertIsInstance(tensor, torch.FloatTensor)
                 self.assertEqual(tensor, torch.FloatTensor([[1.0, 2.0], [3.0, 4.0]]))
 
+    @unittest.skipIf(not torch.cuda.is_available(), 'no CUDA')
+    def test_serialization_map_location_cuda(self):
+        test_file_path = download_file('https://download.pytorch.org/test_data/gpu_tensors.pt')
+
+        def load_bytes():
+            with open(test_file_path, 'rb') as f:
+                data = io.BytesIO(f.read())
+            return data
+
+        fileobject_lambdas = [lambda: test_file_path, load_bytes]
+        map_locations = [{'cuda:0': 'cuda:0'}, 'cuda:0', torch.device('cuda'), torch.device('cuda', 0)]
+
+        for fileobject_lambda in fileobject_lambdas:
+            for map_location in map_locations:
+                tensor = torch.load(fileobject_lambda(), map_location=map_location)
+                self.assertIsInstance(tensor, torch.cuda.FloatTensor)
+                self.assertEqual(tensor, torch.cuda.FloatTensor([[1.0, 2.0], [3.0, 4.0]]))
+
     def test_serialization_filelike_api_requirements(self):
         filemock = FilelikeMock(b'', has_readinto=False)
         tensor = torch.randn(3, 5)
diff --git a/torch/serialization.py b/torch/serialization.py
@@ -273,15 +273,15 @@ def load(f, map_location=None, pickle_module=pickle):
     Args:
         f: a file-like object (has to implement read, readline, tell, and seek),
             or a string containing a file name
-        map_location: a function, string or a dict specifying how to remap storage
+        map_location: a function, torch.device, string or a dict specifying how to remap storage
             locations
         pickle_module: module used for unpickling metadata and objects (has to
             match the pickle_module used to serialize file)
 
     Example:
         >>> torch.load('tensors.pt')
         # Load all tensors onto the CPU
-        >>> torch.load('tensors.pt', map_location='cpu')
+        >>> torch.load('tensors.pt', map_location=torch.device('cpu'))
         # Load all tensors onto the CPU, using a function
         >>> torch.load('tensors.pt', map_location=lambda storage, loc: storage)
         # Load all tensors onto GPU 1
@@ -318,6 +318,16 @@ def restore_location(storage, location):
     elif isinstance(map_location, _string_classes):
         def restore_location(storage, location):
             return default_restore_location(storage, map_location)
+    elif isinstance(map_location, torch.device):
+        if map_location.type == 'cpu':
+            map_str = 'cpu'
+        elif map_location.type == 'cuda':
+            map_str = 'cuda:{}'.format(map_location.index or 0)
+        else:
+            raise ValueError("The given map_location device is not a cpu or cuda")
+
+        def restore_location(storage, location):
+            return default_restore_location(storage, map_str)
     else:
         def restore_location(storage, location):
             result = map_location(storage, location)
