[BugFix] resilient _exclude_td_from_pytree

ghstack-source-id: 7b3ee829689779777d301f0cfff119e48567f9bb Pull Request resolved: #1038
pytorch · Oct 11, 2024 · 49d226c · 49d226c
1 parent 1e32195
commit 49d226c
Showing 1 changed file with 6 additions and 3 deletions.
diff --git a/tensordict/nn/functional_modules.py b/tensordict/nn/functional_modules.py
@@ -128,11 +128,14 @@ def __init__(self):
 
     def __enter__(self):
         for tdtype in PYTREE_REGISTERED_TDS + PYTREE_REGISTERED_LAZY_TDS:
-            self.tdnodes[tdtype] = SUPPORTED_NODES.pop(tdtype)
+            node = SUPPORTED_NODES.pop(tdtype, None)
+            if node is None:
+                continue
+            self.tdnodes[tdtype] = node
 
     def __exit__(self, exc_type, exc_val, exc_tb):
-        for tdtype in PYTREE_REGISTERED_TDS + PYTREE_REGISTERED_LAZY_TDS:
-            SUPPORTED_NODES[tdtype] = self.tdnodes[tdtype]
+        for tdtype, node in self.tdnodes.items():
+            SUPPORTED_NODES[tdtype] = node
 
     def set(self):
         self.__enter__()