I think I got it. It does not seem like a very clean solution, but it works
param_groups = [{"params": [latent[i].detach().requires_grad_(True)]} for i in range(latent.shape[0])] 
opt = optim.Adam(params=param_groups)

# in my loop
generated = self.do_some_fancy_stuff(latent)
loss = my_loss(generated, groundtruth)

# Do zero_grad manually
for i in range(latent.shape[0]):
    opt.param_groups[i]["params"][0].grad = None
latent.grad = None if latent.grad is None else latent.grad.detach().zero_()

loss.backward()

# Copy the gradients from the slices back to the original tensor
for i in range(latent.shape[0]):
    opt.param_groups[i]["params"][0].grad = latent.grad[i]

# Apply the gradients to the latent vectors
opt.step()

# Copy the new latent vectors from the slices to the original tensor so that it can be used in the next epoch
for i in range(latent.shape[0]):
    latent.data[i] = opt.param_groups[i]["params"][0].data

So basically I have an optimizer with one group for each sample. In my backprop loop I run my model on the original tensor (not the sliced), calculate the loss and gradient, copy the gradient to the slices, run the optimizer and copy the new values from the slices back to the original tensor.
This seems to work, but is this the best way to do it? The performance is not ideal because I am copying data a lot of times. It would be best if there is a way to tell PyTorch that the slices and the original tensor should share the same memory and gradients. Is this possible?
Of course, the LR scheduler still needs to be added now, but that should be not a big problem.
EDIT:
Ok, seems like the LR scheduler is updating each parameter group at the same time. So I will need a separate optimizer and scheduler for each sample. This seems like it is getting a really dirty solution. Am I on the wrong track?