Repository: pytorch/tensordict (language: python)

Issue:
[Test] fix inline TDParams kwargs for nontensordata
Stack from [ghstack](https://github.com/ezyang/ghstack) (oldest at bottom):
* __->__ #1095




Discussion / hints:

Fix patch:
diff --git a/tensordict/nn/probabilistic.py b/tensordict/nn/probabilistic.py
index 4567d65c5..3b254731d 100644
--- a/tensordict/nn/probabilistic.py
+++ b/tensordict/nn/probabilistic.py
@@ -20,6 +20,7 @@
 from tensordict.nn.sequence import TensorDictSequential
 
 from tensordict.nn.utils import _set_skip_existing_None
+from tensordict.tensorclass import is_non_tensor
 from tensordict.tensordict import TensorDictBase
 from tensordict.utils import _zip_strict
 from torch import distributions as D, Tensor
@@ -638,5 +639,9 @@ def _dynamo_friendly_to_dict(data):
         return data
     if isinstance(data, TensorDictBase):
         # to_dict is recursive and we don't want that
-        return {data[key] for key in data.keys()}
+        items = dict(data.items())
+        for k, v in items.items():
+            if is_non_tensor(v):
+                items[k] = v.data
+        return items
     return data


Test patch:
diff --git a/test/test_compile.py b/test/test_compile.py
index d714618d2..ead3c10e6 100644
--- a/test/test_compile.py
+++ b/test/test_compile.py
@@ -16,6 +16,7 @@
 
 from tensordict import (
     assert_close,
+    NonTensorData,
     PYTREE_REGISTERED_LAZY_TDS,
     PYTREE_REGISTERED_TDS,
     tensorclass,
@@ -665,7 +666,9 @@ def test_dispatch_tensor(self, mode):
         torch.testing.assert_close(mod(x=x, y=y), mod_compile(x=x, y=y))
 
     def test_prob_module_with_kwargs(self, mode):
-        kwargs = TensorDictParams(TensorDict(scale=1.0), no_convert=True)
+        kwargs = TensorDictParams(
+            TensorDict(scale=1.0, validate_args=NonTensorData(False)), no_convert=True
+        )
         dist_cls = torch.distributions.Normal
         mod = Mod(torch.nn.Linear(3, 3), in_keys=["inp"], out_keys=["loc"])
         prob_mod = Seq(