Hi,
Indeed, you will need to keep the nn.Parameter on the CPU and only perform computation on the GPU in a differentiable manner. That way, it will automatically accumulate back onto the CPU parameters.
Something like:
full_params = torch.rand(num_classes, feature_size, device="cpu")

loss = 0.
for slice in slices:
  param_subset = full_params[slice]
  param_subset = param_subset.cuda()
  # Compute that part of the loss
  loss += local_loss
loss.backward()

Does that match what you want?