joblib · landoskape · Jun 28, 2024 · Jun 28, 2024 · Jun 28, 2024 · Jun 28, 2024
diff --git a/joblib/hashing.py b/joblib/hashing.py
@@ -241,9 +241,54 @@
         Hasher.save(self, obj)
 
 
+class TorchHasher(NumpyHasher):
+    """ Special case for the hasher for when torch is loaded.
+
+        This class extends the NumpyHasher class to handle torch tensors and
+        torch modules. It converts torch tensors and torch modules to numpy
+        arrays for deterministic hashing.
+    """
+
+    def __init__(self, hash_name="md5", coerce_mmap=False):
+        super().__init__(hash_name, coerce_mmap)
+        from torch.nn import Module as torch_nnModule  # noqa: import-outside-toplevel
+        from torch import Tensor as torch_Tensor  # noqa: import-outside-toplevel
+
+        self.torch_nnModule = torch_nnModule
+        self.torch_Tensor = torch_Tensor
+
+    def _convert_tensors_to_numpy(self, obj):
+        # Recursively convert torch tensors in obj to numpy arrays
+        if isinstance(obj, dict):
+            for key, value in obj.items():
+                obj[key] = self._convert_tensors_to_numpy(value)
+        if isinstance(obj, self.torch_nnModule):
+            state_dict = obj.state_dict()
+            obj = {key: self._convert_tensors_to_numpy(value)
+                   for key, value in state_dict.items()}
+            return obj
+        if isinstance(obj, self.torch_Tensor):
+            obj_as_numpy = obj.cpu().detach().numpy()
+            return obj_as_numpy
+        return obj
+
+    def save(self, obj):
+        """ Subclass again to convert torch tensors and torch modules to numpy
+            arrays for deterministic hashing. Torch tensors do not have
+            deterministic pickle representations and therefore hashing them is
+            not reliable.
+
+            Torch tensors are converted to numpy arrays directly, and torch
+            modules are converted to dictionaries of numpy arrays
+            corresponding to each component in the state_dict.
+        """
+        obj = self._convert_tensors_to_numpy(obj)
+        NumpyHasher.save(self, obj)
+
+
 def hash(obj, hash_name='md5', coerce_mmap=False):
     """ Quick calculation of a hash to identify uniquely Python objects
-        containing numpy arrays.
+        containing numpy arrays or torch tensors.
 
         Parameters
         ----------
@@ -258,7 +303,9 @@
         raise ValueError("Valid options for 'hash_name' are {}. "
                          "Got hash_name={!r} instead."
                          .format(valid_hash_names, hash_name))
-    if 'numpy' in sys.modules:
+    if 'torch' in sys.modules:
+        hasher = TorchHasher(hash_name=hash_name, coerce_mmap=coerce_mmap)
+    elif 'numpy' in sys.modules:
         hasher = NumpyHasher(hash_name=hash_name, coerce_mmap=coerce_mmap)
     else:
         hasher = Hasher(hash_name=hash_name)