Hi @albanD ,
If I will just create a tensor how will that tensor be moved to GPU?
I don’t want to use .to(device) because that won’t work with DataParallel.
Are you saying that all the tensors defined under autograd.Function will be automatically taken to device when I will do model.cuda() or torch.nn.DataParallel(model)?
Example:
class Myclass(Function):
@staticmethod
def forward(ctx,input):
mytensor = torch.tensor([1])
So a model defined using above class (Myclass) when taken to device (by model.cuda() or torch.nn.DataParallel(model)) will take all the tensors in forward to the device?