Implement the get_device method in the storage base class.

pytorch · Apr 23, 2023 · cf589af · cf589af
1 parent 7a8d0cc
commit cf589af
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/torch/storage.py b/torch/storage.py
@@ -39,7 +39,10 @@ def size(self) -> int:
     def type(self, dtype: str = None, non_blocking: bool = False) -> T: ...  # noqa: E704
     def cuda(self, device=None, non_blocking=False, **kwargs) -> T: ...  # noqa: E704
     def element_size(self) -> int: ...  # noqa: E704
-    def get_device(self) -> int: ...  # noqa: E704
+
+    def get_device(self) -> int:
+        return self.device.index
+
     def data_ptr(self) -> int: ...  # noqa: E704
 
     # Defined in torch/csrc/generic/StorageSharing.cpp