amend

pytorch · Jun 18, 2024 · 0a8000b · 0a8000b
1 parent fc1ec76
commit 0a8000b
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/torchrl/envs/libs/jax_utils.py b/torchrl/envs/libs/jax_utils.py
@@ -129,7 +129,7 @@ def _tensordict_to_object(tensordict: TensorDictBase, object_example):
             # We need to flatten to fix https://github.com/pytorch/rl/issues/2184
             value = value.contiguous()
             value = value.detach()
-            if value.ndim:
+            if value.ndim > 1:
                 value = value.flatten().clone()
             else:
                 value = value.clone()