[BUG] VLAWrapperBase keeps constructor self.device

[YeonwooSung]

Describe the bug

VLAWrapperBase is an nn.Module and stores the constructor device as Python state:

self.device = None if device is None else torch.device(device)

The kwarg is not documented. Runtime placement already uses out.device / chunk.device, so device="cuda" is a silent no-op for the base class: it neither moves parameters nor is consulted after init.

.to("meta") / .cpu() / .cuda() move parameters and buffers, not that Python attribute. A CPU-constructed wrapper still reports self.device == torch.device("cpu") after .to("meta").

Related to #4224, #4345, and #4346.

To Reproduce

import torch
from tensordict import NonTensorStack, TensorDict
from torchrl.modules.vla import TinyVLA, VLAWrapperBase

base = VLAWrapperBase(action_dim=2, chunk_size=2, device="cpu")
assert base.device == torch.device("cpu")
base.to("meta")
assert base.device == torch.device("cpu")  # stale: .to() did not update it

policy = TinyVLA(action_dim=2, chunk_size=2, device="cpu")
td = TensorDict(
    {
        "observation": {
            "image": torch.zeros(1, 3, 16, 16, dtype=torch.uint8),
            "state": torch.zeros(1, 5),
        },
        "language_instruction": NonTensorStack("pick"),
    },
    batch_size=[1],
)
policy(td)  # materialize lazy modules
policy.to("meta")
assert next(policy.parameters()).device.type == "meta"
assert policy.device == torch.device("cpu")  # stale constructor cache

Expected behavior

Do not keep a single-device attribute on the module. The constructor may still accept device= so subclasses can place parameters at init (TinyVLA already calls self.to(device)). It must not retain that value as module state. After .to("meta"), there must be no Python self.device left at the constructor value. Runtime placement should keep using the involved tensor (chunk.device / out.device).

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)

[YeonwooSung]

/assign