[BugFix] Fix error in state_dict tests (#531)

pytorch · Sep 20, 2023 · 785d11f · 785d11f · github-actions · Sep 20, 2023
1 parent a721766
commit 785d11f
Show file tree

Hide file tree

Showing 2 changed files with 1 addition and 3 deletions.
diff --git a/tensordict/tensordict.py b/tensordict/tensordict.py
@@ -3456,7 +3456,6 @@ def unflatten_keys(self, separator: str = ".", inplace: bool = False) -> T:
             if key in keys and (
                 not is_tensor_collection(out.get(key)) or not out.get(key).is_empty()
             ):
-                print(out.get(key))
                 raise KeyError(
                     "Unflattening key(s) in tensordict will override existing unflattened key"
                 )

diff --git a/test/test_nn.py b/test/test_nn.py
@@ -3172,7 +3172,6 @@ def test_sd_params(self, detach):
                 k: v if not isinstance(v, torch.Tensor) else v * 0
                 for k, v in sd.items()
             }
-            print(sd)
         # do some op to create a graph
         td.apply(lambda x: x + 1)
         # load the data
@@ -3181,7 +3180,7 @@ def test_sd_params(self, detach):
         assert (td == 0).all()
 
     def test_sd_module(self):
-        td = TensorDict({"1": 1, "2": 2, "3": {"3": 3}}, [])
+        td = TensorDict({"1": 1.0, "2": 2.0, "3": {"3": 3.0}}, [])
         td = TensorDictParams(td)
         module = nn.Linear(3, 4)
         module.td = td