Idiomatic way to place custom tensors on GPU

Hi, what’s the idiomatic PyTorch way of putting tensor fields on custom models on the GPU? An example should clarify what I’m asking.

Say I have a class like this:

class MyModel(nn.Module):
  def __init__(self):
    # ...
    self.my_tensor = torch.rand(...)

  def forward(self, x):
    # ...
    out = out + self.my_tensor
    # ...

# ...
model =
model.my_tensor = model.my_tensor = to(device)

Wondering if there’s a better way to put all custom tensors on device without manually needing to run .to(device) on each.

All nn.Parameters and buffers (registered via self.register_buffer) will automatically be pushed to the specified device when calling The difference would be that the former are trainable (requires_grad is set to True by default) while the latter do not. In case you don’t want to store the buffer in the state_dict, you could use self.register_buffer('name', tensor, persistent=False).

1 Like

Thanks! That’s exactly the kind of thing I was looking for.