Autograd will make sure to track all to(), cuda(), and cpu() operations.
Here is a small example:
class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.emb = nn.Embedding(10, 100)
        self.fc = nn.Linear(100, 2)
        
    def forward(self, x):
        x = self.emb(x)
        x = x.to(self.fc.weight.device)
        x = self.fc(x)
        return x

model = MyModel()
model.fc.cuda()

x = torch.randint(0, 10, (20,))
out = model(x)
print(out.device)
> cuda:0
out.mean().backward()

for name, param in model.named_parameters():
    print('{}, param.grad.abs().sum() {}, param.device {}, grad.device {}'.format(
        name, param.grad.abs().sum(), param.device, param.grad.device))

> emb.weight, param.grad.abs().sum() 3.2752888202667236, param.device cpu, grad.device cpu
fc.weight, param.grad.abs().sum() 31.677473068237305, param.device cuda:0, grad.device cuda:0
fc.bias, param.grad.abs().sum() 1.0, param.device cuda:0, grad.device cuda:0