[BugFix] Fix tensordict.get in TDModule tensor retrieval for NonTensorData

ghstack-source-id: 240b54a
Pull Request resolved: #1249

Author: Vincent Moens

tensordict/nn/common.py

@@ -1169,7 +1169,9 @@ def forward(
             if self._kwargs is not None:
                 kwargs.update(
                     {
-                        kwarg: tensordict.get(in_key, default=default)
+                        kwarg: tensordict._get_tuple_maybe_non_tensor(
+                            _unravel_key_to_tuple(in_key), default=default
+                        )
                         for kwarg, in_key in _zip_strict(self._kwargs, self.in_keys)
                     }
                 )

test/test_nn.py

@@ -689,6 +689,24 @@ def test_nontensor(self):
             out_keys=["out"],
         )
         assert tdm(TensorDict())["out"] == "a string!"
+        tdm = TensorDictModule(
+            lambda a_string: a_string + " is a string!",
+            in_keys=["string"],
+            out_keys=["another string"],
+        )
+        assert (
+            tdm(TensorDict(string="a string"))["another string"]
+            == "a string is a string!"
+        )
+        tdm = TensorDictModule(
+            lambda string: string + " is a string!",
+            in_keys={"string": "key"},
+            out_keys=["another string"],
+            out_to_in_map=True,
+        )
+        assert (
+            tdm(TensorDict(key="a string"))["another string"] == "a string is a string!"
+        )
 
     @pytest.mark.parametrize(
         "out_keys",