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.