No, this won’t be possible as slicing a tensor will create a non-leaf tensor. You could create the sliced tensor with requires_grad=True, pass it to the optimizer, and concatenate or stack it with the other tensor before using it in the actual operation. Alternatively, you could also try to zero out the gradients of the parts of the tensor which should not be updated.