Ah, I see.  When you write something like XYZ[0] = ... , this is still an in-place operation.
My recommendation to follow a scheme of conditionally changing values in your loss function (I assume from a network output of XYZ) is to initialize new tensors for each index:
Here’s an example for a tensor XYZ of shape [3]:
XYZ = nn.Parameter(torch.tensor([.008,.009,1.0]))

if XYZ[0] > 0.008856:
    at_0 = XYZ[0] ** 1/3
else:
    at_0 = (7.787 * XYZ[0]) + (16 / 116)

if XYZ[1] > 0.008856:
    at_1 = XYZ[1] ** 1/3
else:
    at_1 = (7.787 * XYZ[1]) + (16 / 116)
    
at_2 = XYZ[2]

test = torch.stack ([at_0, at_1, at_2])
test.sum().backward()

The tensor slicing will be kept track of automatically by autograd too:
print (XYZ[1])
tensor(0.0090, grad_fn=<SelectBackward>)