Update

[ghstack-poisoned]
pytorch · Nov 29, 2024 · e851a5c · e851a5c
1 parent 2a6262e
commit e851a5c
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 2 deletions.
diff --git a/tensordict/nn/probabilistic.py b/tensordict/nn/probabilistic.py
@@ -1119,7 +1119,18 @@ def forward(
         else:
             result = tensordict_exec
             if self._select_before_return:
-                return tensordict.update(result, keys_to_update=self.out_keys)
+                # We must also update any value that has been updated during the course of execution
+                # from the input data.
+                if is_compiling():
+                    keys = [  # noqa: C416
+                        k
+                        for k in {k for k in self.out_keys}.union(  # noqa: C416
+                            {k for k in tensordict.keys(True, True)}  # noqa: C416
+                        )
+                    ]
+                else:
+                    keys = list(set(self.out_keys + list(tensordict.keys(True, True))))
+                return tensordict.update(result, keys_to_update=keys)
         return result
 
 

diff --git a/tensordict/nn/sequence.py b/tensordict/nn/sequence.py
@@ -35,6 +35,10 @@
     )
     FUNCTORCH_ERROR = "functorch not installed. Consider installing functorch to use this functionality."
 
+try:
+    from torch.compiler import is_compiling
+except ImportError:
+    from torch._dynamo import is_compiling
 
 __all__ = ["TensorDictSequential"]
 
@@ -491,7 +495,15 @@ def forward(
             if self._select_before_return:
                 # We must also update any value that has been updated during the course of execution
                 # from the input data.
-                keys = list(set(self.out_keys + list(tensordict.keys(True, True))))
+                if is_compiling():
+                    keys = [  # noqa: C416
+                        k
+                        for k in {k for k in self.out_keys}.union(  # noqa: C416
+                            {k for k in tensordict.keys(True, True)}  # noqa: C416
+                        )
+                    ]
+                else:
+                    keys = list(set(self.out_keys + list(tensordict.keys(True, True))))
                 return tensordict.update(result, keys_to_update=keys)
         return result