Update

[ghstack-poisoned]
pytorch · Sep 25, 2024 · f2e6dd3 · f2e6dd3
1 parent 7549295
commit f2e6dd3
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tensordict/nn/functional_modules.py b/tensordict/nn/functional_modules.py
@@ -311,7 +311,7 @@ def incompatible_error():
             flat_outputs.append(out)
         return tree_unflatten(flat_outputs, output_spec)
 
-        vmap_src._unwrap_batched = _unwrap_batched
+    vmap_src._unwrap_batched = _unwrap_batched
 
 
 # Tensordict-compatible Functional modules