Where does a new tensor go

def add_noise(x):
    return x + torch.randn(x.shape)

Fine on a laptop. On a GPU:

RuntimeError: Expected all tensors to be on the same device,
but found at least two devices, cuda:0 and cpu!

torch.randn creates on the default device, which is the CPU, and the function has no idea x is anywhere else.

Three ways to say “wherever x is”

torch.randn(x.shape, device=x.device, dtype=x.dtype)
torch.randn_like(x)
x.new_empty(x.shape).normal_()

The _like family is the shortest and copies both the device and the dtype, which is usually what you want and is easy to forget with the explicit form. new_* methods do the same and let you change the shape.

Why this is a habit and not a special case

Any function that creates a tensor and combines it with an argument has this bug latent in it: masks, positional encodings, causal triangles, noise, constants built with torch.tensor(...), index ranges from torch.arange.

Writing device-agnostic code is not about supporting exotic hardware. It is what makes a module you wrote on CPU work inside a model someone moves to a GPU with model.cuda(), which moves parameters and buffers and cannot reach inside your forward pass.

The related trap

A constant tensor created in __init__ and stored as a plain attribute is not moved by model.to(device) either, because it is not a parameter or a buffer. register_buffer is the fix, and the modules track returns to it.

Your task

def add_scaled_range(x: torch.Tensor, scale: float) -> torch.Tensor

Return x + scale * arange(n), where n is the length of x, with the range created on the same device and dtype as x.

The starter creates it with the defaults.