Could you check the shape of kernels and the value of batch_size?
If I’m not mistaken this code would only work, if the shapes are batch_size*2 as seen here:
# works
kernels = torch.randn(128, 128)
batch_size = 64

XX = kernels[:batch_size, :batch_size]
YY = kernels[batch_size:, batch_size:]
XY = kernels[:batch_size, batch_size:]
YX = kernels[batch_size:, :batch_size]
loss = torch.mean(XX + YY - XY - YX)

# fails
kernels = torch.randn(128, 127)
batch_size = 64

XX = kernels[:batch_size, :batch_size]
YY = kernels[batch_size:, batch_size:]
XY = kernels[:batch_size, batch_size:]
YX = kernels[batch_size:, :batch_size]
loss = torch.mean(XX + YY - XY - YX)