Repository: pytorch/tensordict

Issue:
[BUG] TensorDict does not properly handle nested TensorClasses
## Describe the bug

When placing a TensorClass inside a TensorDict, the `to`-operations breaks
Ie  `to_tensorclass` is called on the TensorClasses, converting it to a TensorDict which has lost its non-tensor attributes. 

## To Reproduce

Run the following code:

```python
import torch
from tensordict import tensorclass, TensorDict


@tensorclass
class TensorClass:
    tens: torch.Tensor
    order: tuple[str]
    test: str

empty_batchsize = torch.Size([])

def main():
    c = TensorClass(torch.ones(1), ("a", "b", "c"), 'Hello',
                    batch_size=empty_batchsize)
    
    sc = torch.stack([c, c, c])

    print("=== TensorClass ===")
    print(sc)
    print("--- idx ---")
    print(sc[0])
    print("--- to ---")
    print(sc.to("cuda:0"))

    td = torch.stack(
        [TensorDict({
            "t": torch.ones(1),
            "c": c
        }, batch_size=empty_batchsize)
        for _ in range(3)])

    print("=== TensorDict ===")
    print(td)
    print("--- idx ---")
    print(td[0])
    print("--- to ---")
    print(td.to("cuda:0"))
    print("--- to td---")
    print(td.to_tensordict())

if __name__ == "__main__":
    main()
```

## Output

```
=== TensorClass ===                                                                                                                                                                                                                                           
TensorClass(                                                                                                                                                                                                                                                  
    tens=Tensor(shape=torch.Size([3, 1]), device=cpu, dtype=torch.float32, is_shared=False),                                                                                                                                                                  
    order=('a', 'b', 'c'),                                                                                                                                                                                                                                    
    test='Hello',                                                                                                                                                                                                                                             
    batch_size=torch.Size([3]),                                                                                                                                                                                                                               
    device=None,                                                                                                                                                                                                                                              
    is_shared=False)                                                                                                                                                                                                                                          
--- idx ---                                                                                                                                                                                                                                                   
TensorClass(                                                                                                                                                                                                                                                  
    tens=Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.float32, is_shared=False),
    order=('a', 'b', 'c'),
    test='Hello',
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)
--- to ---
TensorClass(
    tens=Tensor(shape=torch.Size([3, 1]), device=cuda:0, dtype=torch.float32, is_shared=True),
    order=('a', 'b', 'c'),
    test='Hello',
    batch_size=torch.Size([3]),
    device=cuda:0,
    is_shared=True)
=== TensorDict ===
LazyStackedTensorDict(
    fields={
        c: TensorClass(
            tens=Tensor(shape=torch.Size([3, 1]), device=cpu, dtype=torch.float32, is_shared=False),
            order=('a', 'b', 'c'),
            test='Hello',
            batch_size=torch.Size([3]),
            device=None,
            is_shared=False),
        t: Tensor(shape=torch.Size([3, 1]), device=cpu, dtype=torch.float32, is_shared=False)},
    batch_size=torch.Size([3]),
    device=None,
    is_shared=False)
--- idx ---
TensorDict(
    fields={
        c: TensorClass(
            tens=Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.float32, is_shared=False),
            order=('a', 'b', 'c'),
            test='Hello',
            batch_size=torch.Size([]),
            device=None,
            is_shared=False),
        t: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.float32, is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)
--- to ---
TensorDict(
    fields={
        c: TensorDict(
            fields={
                tens: Tensor(shape=torch.Size([3, 1]), device=cuda:0, dtype=torch.float32, is_shared=True)},
            batch_size=torch.Size([3]),
            device=cuda:0,
            is_shared=True),
        t: Tensor(shape=torch.Size([3, 1]), device=cuda:0, dtype=torch.float32, is_shared=True)},
    batch_size=torch.Size([3]),
    device=cuda:0,
    is_shared=True)
--- to td---
TensorDict(
    fields={
        c: TensorDict(
            fields={
                tens: Tensor(shape=torch.Size([3, 1]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([3]),
            device=None,
            is_shared=False),
        t: Tensor(shape=torch.Size([3, 1]), device=cpu, dtype=torch.float32, is_shared=False)},
    batch_size=torch.Size([3]),
    device=None,
    is_shared=False)
```


## Expected behavior

When calling `to` on the TensorDict, the TensorClass should stay intact, just at is when called directly on the TensorClass.

## System info

Describe the characteristic of your environment:
 * Python 3.9
 * Tensordict 0.1.2
 * Torch 1.13.1
 * Numpy  1.23.5
 * Sys Version, Plattform  3.10.9 (main, Jan 11 2023, 15:21:40) [GCC 11.2.0] linux 

## Reason and Possible fixes

The problem seems to arise when LazyStackedTensorDict is transformed to TensorDict in the `to`-call: https://github.com/pytorch-labs/tensordict/blob/879678435e52d2610e4bb820cb4878cd01caec7e/tensordict/tensordict.py#L6844C1-L6860C22

In it, all values of the TensorDict are transformed to TensorDicts themselves: https://github.com/pytorch-labs/tensordict/blob/879678435e52d2610e4bb820cb4878cd01caec7e/tensordict/tensordict.py#L2365
Instead, TensorClasses should be kept as TensorClasses.
A quick fix is to instead of checking `_is_tensor_collection` in 2368, `issubclass(type(value), TensorDictBase)` should be checked.

Though that keeps one from transforming nested TensorClasses to TensorDicts.

An alternative would be to modify the calling of `to_tensordict` in LazyStackedTensorDicts `to` method.
The idea would be to only convert LazyStackedTensorDicts, not TensorClasses.
Perhaps using by adding an arg to the `to_tensordict` function.

## Checklist

- [x] I have checked that there is no similar issue in the repo (**required**)
- [x] I have read the [documentation](https://github.com/pytorch/rl/tree/main/docs/) (**required**)
- [x] I have provided a minimal working example to reproduce the bug (**required**)



Discussion:
Thanks for this.
I'm on PTO these days so that's why I'm not actively participating in this thread (or others FWIW).
I will catch up as soon as I'm back! Sorry about that.

Resolving patch:
diff --git a/tensordict/tensordict.py b/tensordict/tensordict.py
--- a/tensordict/tensordict.py
+++ b/tensordict/tensordict.py
@@ -6856,7 +6856,12 @@ def to(self, dest: DeviceType | type, **kwargs) -> TensorDictBase:
             dest = torch.device(dest)
             if self.device is not None and dest == self.device:
                 return self
-            td = self.to_tensordict().to(dest, **kwargs)
+            td = LazyStackedTensorDict(
+                *[td.to(dest, **kwargs) for td in self.tensordicts],
+                stack_dim=self.stack_dim,
+                hook_out=self.hook_out,
+                hook_in=self.hook_in,
+            )
             return td
 
         elif isinstance(dest, torch.Size):


Test patch:
diff --git a/test/test_tensorclass.py b/test/test_tensorclass.py
--- a/test/test_tensorclass.py
+++ b/test/test_tensorclass.py
@@ -1672,6 +1672,47 @@ class MyClass:
     assert tc.batch_size == torch.Size([10])
 
 
+class TestNesting:
+    @tensorclass
+    class TensorClass:
+        tens: torch.Tensor
+        order: tuple[str]
+        test: str
+
+    def get_nested(self):
+
+        c = self.TensorClass(torch.ones(1), ("a", "b", "c"), "Hello", batch_size=[])
+
+        td = torch.stack(
+            [TensorDict({"t": torch.ones(1), "c": c}, batch_size=[]) for _ in range(3)]
+        )
+        return td
+
+    def test_to(self):
+        td = self.get_nested()
+        td = td.to("cpu:1")
+        assert isinstance(td.get("c")[0], self.TensorClass)
+
+    def test_idx(self):
+        td = self.get_nested()[0]
+        assert isinstance(td.get("c"), self.TensorClass)
+
+    def test_apply(self):
+        td = self.get_nested()
+        td = td.apply(lambda x: x + 1)
+        assert isinstance(td.get("c")[0], self.TensorClass)
+
+    def test_split(self):
+        td = self.get_nested()
+        td, _ = td.split([2, 1], dim=0)
+        assert isinstance(td.get("c")[0], self.TensorClass)
+
+    def test_chunk(self):
+        td = self.get_nested()
+        td, _ = td.chunk(2, dim=0)
+        assert isinstance(td.get("c")[0], self.TensorClass)
+
+
 if __name__ == "__main__":
     args, unknown = argparse.ArgumentParser().parse_known_args()
     pytest.main([__file__, "--capture", "no", "--exitfirst"] + unknown)