[BUG] CatFrames with padding="same" is slow

[kurtamohler]

Describe the bug

A discord user reported that CatFrames is slow and gave a repro script: https://discord.com/channels/1171857748607115354/1268244652180377732/1270523213600002129

To Reproduce

The given repro script allocated more memory than my computer can handle, so I modified it like so:

import torch
from torchrl.envs import CatFrames, Compose
from torchrl.data import PrioritizedSliceSampler, ReplayBuffer, LazyTensorStorage
from tensordict import TensorDict
import timeit

device = 'cpu'

def catframes_speed_test(with_catframes=True):
    if with_catframes:
        obs_shape = (4, 42, 42, 1)
        transform = Compose(
            CatFrames(N=4, dim=-1, in_keys=["observation"], done_key="done"),
            CatFrames(N=4, dim=-1, in_keys=[("next", "observation")], done_key="done"))
    else:
        obs_shape = (4, 42, 42, 4)
        transform = None


    max_size = 100_000

    sampler = PrioritizedSliceSampler(
        max_capacity=max_size,
        alpha=0.5,
        beta=0.4,
        strict_length=False,
        num_slices=32 // 4,
        span=[True, False],
    )

    exp_buffer = ReplayBuffer(
        storage=LazyTensorStorage(max_size=max_size, device=device),
        sampler=sampler,
        batch_size=32,
        transform=transform,
    )

    fake_data = TensorDict(
        {
            "observation": torch.zeros(obs_shape, dtype=torch.float32, device=device),
            "next": {"observation": torch.zeros(obs_shape, dtype=torch.float32, device=device),
                     "done": torch.zeros((obs_shape[0], 1), dtype=torch.bool, device=device)},
        },
        batch_size=[obs_shape[0]],
        device=device,
    )

    for _ in range(25):
        exp_buffer.extend(fake_data)

    num_iters = 30
    timer = timeit.Timer(
        stmt=lambda: exp_buffer.sample(return_info=True),
    )
    time_per_iter = timer.timeit(num_iters) / num_iters
    print(f"Sampling took {time_per_iter} seconds.")

    exp_buffer.empty()


if __name__ == "__main__":
    print("WITHOUT CATFRAMES")
    catframes_speed_test(with_catframes=False)

    print("WITH CATFRAMES")
    catframes_speed_test(with_catframes=True)

On my machine, I get this output:

WITHOUT CATFRAMES
Sampling took 0.0016307519336502688 seconds.
WITH CATFRAMES
Sampling took 0.23012657026653566 seconds.

Expected behavior

Adding the CatFrames transformation to a replay buffer should not slow it down significantly, but in the above case, it caused a slowdown of more than 100x.

Checklist

• I have checked that there is no similar issue in the repo (required)
• I have read the documentation (required)
• I have provided a minimal working example to reproduce the bug (required)