I was wondering if it was possible to do something like the following, where I try to load the model from CPU -> GPU before the computation and send it back after:
import torch
from torch import nn

DUMMY = torch.empty((), requires_grad=True)
class Clive(torch.autograd.Function):
    @staticmethod
    def forward(ctx, layer, dummy, inputs):
        layer.to(0)
        inputs = inputs
        with torch.enable_grad():
            outputs = layer(inputs)
        ctx.layer = layer
        ctx.inputs = inputs
        ctx.outputs = outputs
        layer.to('cpu')
        return outputs.detach()

    @staticmethod
    def backward(ctx, *grad_outputs):
        ctx.layer.to(0)

        diff_with = ctx.layer.parameters()
        if ctx.inputs.requires_grad:
            diff_with = itertools.chain([ctx.inputs], diff_with)

        grads = torch.autograd.grad(outputs=(ctx.outputs,), inputs=diff_with, grad_outputs=grad_outputs)
        
        if ctx.inputs.requires_grad:
            input_grad, *grads = grads
        else:
            input_grad = None

        for g, l in zip(grads, ctx.layer.parameters()):
            l.backward(g)
        ctx.layer.to('cpu')
        return (None, None, input_grad)

m1 = nn.Linear(10, 10)
x = torch.randn(10, 10, device=0)
Clive.apply(m1, DUMMY, x).sum().backward()
print(m1.weight.grad)

Unfortunately this errors with One of the differentiated Tensors appears to not have been used in the graph. Set allow_unused=True if this is the desired behavior.. I’m still working to understand this error, can anyone dumb it down for me? If I get rid of the .to(...) method calls it works ok.
(I don’t really know if this is practical/useful, I’m just curious if I can actually do it :D)