Hey folks,
I was wondering if someone could clarify some aspects of the cuda context manager.
For instance lets say we have dummy function, and we want to either send the whole function to cuda or whatever tensors are created/manipulated inside that dummy function to automatically get sent to cuda.
def dummy_func(input, output):
    y = input * output
    tmp = torch.ones(input.size(0), output.size(2))
    y.scatter_(1, tmp, 1)

What I would like is everythin inside dummy_func to operate on cuda without having to explicity add to each new tensor .to(device) or .cuda()
What I’ve tried are the following
@torch.cuda()
def dummy_func(...):

and
def dummy_func(....):
    with torch.cuda.device(0):
           .....


But none of those gave the expected outcome.
Is there any way to achieve that in order to avoid having to manually add .cuda to every possible tensor?