amend

pytorch · Jan 18, 2024 · f6d1254 · f6d1254
1 parent 0b9b375
commit f6d1254
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/tensordict/_td.py b/tensordict/_td.py
@@ -846,7 +846,8 @@ def empty():
         def unbind(key, val, tds=tds):
             unbound = (
                 val.unbind(dim)
-                if not _is_tensor_collection(type(val))
+                if not isinstance(val, TensorDictBase)
+                # tensorclass is also unbound using plain unbind
                 else val._unbind(dim)
             )
             for td, _val in zip(tds, unbound):