[BugFix] Fix vmap monkey patching

ghstack-source-id: 69f4795d5cb81db7b79d9c98626414c4cc5ce886 Pull Request resolved: #1009
pytorch · Sep 25, 2024 · eba0769 · eba0769
1 parent d9fece7
commit eba0769
Showing 1 changed file with 0 additions and 9 deletions.
diff --git a/tensordict/nn/functional_modules.py b/tensordict/nn/functional_modules.py
@@ -270,15 +270,6 @@ def _unwrap_batched(
             with _exclude_td_from_pytree():
                 flat_batched_outputs, output_spec = tree_flatten(batched_outputs)
 
-        for out in flat_batched_outputs:
-            # Change here:
-            if isinstance(out, torch.Tensor) or is_tensor_collection(out):
-                continue
-            raise ValueError(
-                f"vmap({_get_name(func)}, ...): `{_get_name(func)}` must only return "
-                f"Tensors, got type {type(out)} as a return."
-            )
-
         def incompatible_error():
             raise ValueError(
                 f"vmap({_get_name(func)}, ..., out_dims={out_dims})(<inputs>): "