[Feature] Better detection of non-tensor data (#685)

(cherry picked from commit 551331d)
pytorch · Mar 25, 2024 · 55c34a7 · 55c34a7
1 parent ab54cc1
commit 55c34a7
Show file tree

Hide file tree

Showing 13 changed files with 596 additions and 163 deletions.
diff --git a/.github/workflows/test-rl-gpu.yml b/.github/workflows/test-rl-gpu.yml
@@ -31,6 +31,7 @@ jobs:
       repository: pytorch/tensordict
       gpu-arch-type: cuda
       gpu-arch-version: ${{ matrix.cuda_arch_version }}
+      timeout: 120
       script: |
         # Set env vars from matrix
         export PYTHON_VERSION=${{ matrix.python_version }}

diff --git a/docs/source/reference/tensorclass.rst b/docs/source/reference/tensorclass.rst
@@ -273,3 +273,4 @@ Here is an example:
 
     tensorclass
     NonTensorData
+    NonTensorStack
diff --git a/tensordict/__init__.py b/tensordict/__init__.py
@@ -16,7 +16,7 @@
 from tensordict.memmap import MemoryMappedTensor
 from tensordict.memmap_deprec import is_memmap, MemmapTensor, set_transfer_ownership
 from tensordict.persistent import PersistentTensorDict
-from tensordict.tensorclass import NonTensorData, tensorclass
+from tensordict.tensorclass import NonTensorData, NonTensorStack, tensorclass
 from tensordict.utils import (
     assert_allclose_td,
     is_batchedtensor,
@@ -43,6 +43,7 @@
     "TensorDict",
     "TensorDictBase",
     "merge_tensordicts",
+    "NonTensorStack",
     "set_transfer_ownership",
     "pad_sequence",
     "is_memmap",