From 89a8a69cfdc34cbe46ba6070a1471e4ca80e0e17 Mon Sep 17 00:00:00 2001
From: Vincent Moens <vmoens@meta.com>
Date: Wed, 20 Nov 2024 09:32:44 +0000
Subject: [PATCH] Update

[ghstack-poisoned]
---
 tensordict/nn/probabilistic.py | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

diff --git a/tensordict/nn/probabilistic.py b/tensordict/nn/probabilistic.py
index 2f344d8a6..4567d65c5 100644
--- a/tensordict/nn/probabilistic.py
+++ b/tensordict/nn/probabilistic.py
@@ -637,8 +637,6 @@ def _dynamo_friendly_to_dict(data):
     if not is_compiling():
         return data
     if isinstance(data, TensorDictBase):
-        items = list(data.items())
-        if not items:
-            return {}
-        return dict(items)
+        # to_dict is recursive and we don't want that
+        return {data[key] for key in data.keys()}
     return data