The constant that stayed behind

A module holding a constant it needs in forward: a causal mask, a set of frequencies, a normalisation scale.

class Layer(nn.Module):
    def __init__(self, n):
        super().__init__()
        self.w = nn.Parameter(torch.ones(n))
        self.scale = torch.full((n,), 2.0)     # a plain attribute

    def forward(self, x):
        return x * self.w * self.scale

Works. Then:

layer.cuda()
layer(x)      # RuntimeError: expected all tensors on the same device

What .to() actually reaches

Module.to() walks the module’s parameters and buffers and moves each one. A plain tensor attribute is neither, so it is invisible to the walk and stays exactly where it was created.

The fix is to tell the module the tensor exists:

self.register_buffer("scale", torch.full((n,), 2.0))

A buffer is state that moves with the module and is saved in state_dict, but is not a parameter: no gradient, and no optimiser will update it. Running statistics in BatchNorm, causal masks, and RoPE frequency tables are all buffers.

The same walk drives everything else

.to(), .cuda(), .half(), state_dict(), load_state_dict() and parameters() all traverse the same registries. A tensor that is not registered is absent from every one of them, which means the device bug arrives with a silent twin: the constant is missing from the checkpoint, and reloading rebuilds it from __init__ instead. If it was ever computed from data, it is now wrong and nothing said so.

When you do not want it in the checkpoint

self.register_buffer("mask", mask, persistent=False)

Moves with the module, stays out of state_dict. Right for anything derivable from the config, like a causal mask, where saving it wastes space and hard-codes a sequence length into the checkpoint.

Your task

Fix Layer so its scale moves with the module, then:

def scale_dtype_after_cast(n: int) -> dict

builds a Layer(n), calls .to(torch.float64) on it, and reports the dtype of w and of scale, plus whether scale appears in state_dict.