[Refactor] Graceful as_tensor (#549)

pytorch · Nov 2, 2023 · 09b327c · 09b327c · github-actions · Nov 2, 2023
1 parent c3caa76
commit 09b327c
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 8 deletions.
diff --git a/tensordict/tensordict.py b/tensordict/tensordict.py
@@ -1811,13 +1811,14 @@ def as_tensor(self):
         and will raise an exception in all other cases.
 
         """
-        try:
-            return self._fast_apply(lambda x: x.as_tensor())
-        except AttributeError as err:
-            raise AttributeError(
-                f"{self.__class__.__name__} does not have an 'as_tensor' method "
-                f"because at least one of its tensors does not support this method."
-            ) from err
+
+        def as_tensor(x):
+            try:
+                return x.as_tensor()
+            except AttributeError:
+                return x
+
+        return self._fast_apply(as_tensor)
 
     def update(
         self,

diff --git a/tensordict/utils.py b/tensordict/utils.py
@@ -767,7 +767,6 @@ def _is_shared(tensor: torch.Tensor) -> bool:
     elif isinstance(tensor, KeyedJaggedTensor):
         return False
     else:
-        print(type(tensor))
         return tensor.is_shared()