[BugFix] Fix output of SipHash(as_tensor=False)

pytorch · Dec 18, 2024 · 26173a4 · 26173a4
1 parent 91064bc
commit 26173a4
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 1 deletion.
diff --git a/test/test_storage_map.py b/test/test_storage_map.py
@@ -46,6 +46,15 @@ def test_sip_hash(self):
         hash_b = torch.tensor(hash_module(b))
         assert (hash_a == hash_b).all()
 
+    def test_sip_hash_nontensor(self):
+        a = torch.rand((3, 2))
+        b = a.clone()
+        hash_module = SipHash(as_tensor=False)
+        hash_a = hash_module(a)
+        hash_b = hash_module(b)
+        assert len(hash_a) == 3
+        assert hash_a == hash_b
+
     @pytest.mark.parametrize("n_components", [None, 14])
     @pytest.mark.parametrize("scale", [0.001, 0.01, 1, 100, 1000])
     def test_randomprojection_hash(self, n_components, scale):

diff --git a/torchrl/data/map/hash.py b/torchrl/data/map/hash.py
@@ -111,7 +111,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor | List[bytes]:
             hash_value = x_i.tobytes()
             hash_values.append(hash_value)
         if not self.as_tensor:
-            return hash_value
+            return hash_values
         result = torch.tensor([hash(x) for x in hash_values], dtype=torch.int64)
         return result