Looking at Pytorch’s Adam implementation, I see this:
for group in self.param_groups:
    for p in group['params']:
        if p.grad is None:
            continue

Unfortunately, this will not look at the gradient of each slice of the parameter, but at the gradient of the parameter as a whole. So I would need to splice the samples and initialize Adam like this:
adam = torch.optim.Adam(params=[samples[i] for i in range(512)])

This unfortunately leads to a new problem. The slices are no leafs anymore and when I try to use this instead, the original samplesis not updated anymore, but just the slices:
adam = torch.optim.Adam(params=[samples[i].detach().requires_grad_(True) for i in range(512)])

Any idea how to solve the issue elegantly?