Idiomatic way to place custom tensors on GPU

1

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.to(device)
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.

2

All nn.Parameters and buffers (registered via self.register_buffer) will automatically be pushed to the specified device when calling model.to(device). 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
3

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