Actually it doesn’t seem to work for my use case - I have to reload the model to GPU before running the backward pass. I’ve mimicked a trivial training loop here:
import torch
from torch import nn

m = nn.Linear(10, 10)
inputs = [
    torch.randn(10, 10, device=0),
    torch.randn(10, 10, device=0)
]
for input in inputs:
    m.to(0)
    activation = m(input)
    m.to('cpu')
    loss = activation.sum()
    # m.to(0) <--- without this copy, it doesn't work
    loss.backward()    
    # m.to('cpu')
    m.zero_grad()

Obviously in this example this is fine, but what I’m aiming to do is only load the layer’s weight into GPU memory when they are needed (e.g. forward and backward). That was what I was going for with the autograd Function.