I was also interested in offloading the computed outputs to the CPU (like gradient checkpointing but not recomputing), but you’re right I can just wrap the model.
Speaking of, is there something obvious I am doing wrong here? It seems to fail if I try to run another backward pass when there’s already gradients allocated.
import torch
from torch import nn
class Wrapped(nn.Module):
    def __init__(self):
        super().__init__()
        self.l = nn.Linear(10, 10)
        
    def forward(self, x):
        self.l.to(0)
        ans = self.l(x)
        self.l.to('cpu')
        return ans

inp = torch.randn(10, 10, device=0)
w = Wrapped()

# Standard forward/backward
w(inp).sum().backward()

# Clear grads and do it again
for p in w.parameters():
    p.grad = None
w(inp).sum().backward()

# Zero grads like we should do, fails with "Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!"
w.zero_grad()
w(inp).sum().backward()