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.