diff --git a/test/test_storage_map.py b/test/test_storage_map.py index db2d0bc2c49..90b16db00d3 100644 --- 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 index 59526628dbe..a3ae9ec1ae9 100644 --- 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