[BugFix] Initialising the classes LazyTensorStorage with a nested TensorDict raises error

[albertbou92]

Description

Initialising the class LazyTensorStorage with a nested TensorDict raises the following error

RuntimeError: the number of sizes provided (1) must be greater or equal to the number of dimensions in the TensorDict (2)

It it due to a call to tensor.expand . I just added the the dimensions of tensor to the output shape, and now it does not crash anymore.

The error can be reproduced with the following code:

import torch
from tensordict import TensorDict
from torchrl.data.replay_buffers.storages import LazyTensorStorage

dummy_td = TensorDict({
    "key1": torch.ones(3, 4),
    "key2": torch.ones(3, 4),
    "next": TensorDict({
        "next_key1": torch.ones(3, 4),
        "next_key2": torch.ones(3, 4),
    }, [3, 4])
}, [3, 4])

mystorage = LazyTensorStorage(1000)
mystorage._init(dummy_td)

Types of changes

What types of changes does your code introduce? Remove all that do not apply:

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds core functionality)
• Breaking change (fix or feature that would cause existing functionality to change)
• Documentation (update in the documentation)
• Example (update in the folder of examples)

Checklist

Go over all the following points, and put an x in all the boxes that apply.
If you are unsure about any of these, don't hesitate to ask. We are here to help!

• I have read the CONTRIBUTION guide (required)
• My change requires a change to the documentation.
• I have updated the tests accordingly (required for a bug fix or a new feature).
• I have updated the documentation accordingly.

[vmoens]

Good catch!
can you look at LazyMemmapStorage, where this error has been solved? You can pretty much copy-paste the _init method from there

[vmoens]

I might be missing something but wasn't the implementation of lazymemmaptensor already solving that issue?
What bug did you get with it when creating a nested TD?
Since it's what we do in the training scripts I would have expected the tests of the training scripts to fail.

[vmoens]

I might be missing something but wasn't the implementation of lazymemmaptensor already solving that issue?
What bug did you get with it when creating a nested TD?
Since it's what we do in the training scripts I would have expected the tests of the training scripts to fail.

[albertbou92]

I might be missing something but wasn't the implementation of lazymemmaptensor already solving that issue? What bug did you get with it when creating a nested TD? Since it's what we do in the training scripts I would have expected the tests of the training scripts to fail.

Seems like I did not have the latest version of tensordict, that is why I was getting also an error for LazyTensorStorage. I reverted the latest changes.

And then fixed LazyTensorStorage _init based on the solution in LazyMemmapStorage _init

[vmoens]

LGTM
Do you mind copying the code snippet that caused the error in a test in test_rb.py for the tensor and memmap version?

[albertbou92]

LGTM Do you mind copying the code snippet that caused the error in a test in test_rb.py for the tensor and memmap version?

done!