Possible wrong assignment in `contrast_matrix` function

[GladerJ]

Hello, thank you for sharing the great code.

In the contrast_matrix function (src/train_plus_utils.py), there seems to be a mistake around line ~218.

matrix_02 = torch.zeros((245, 4))
for i in range(4):
    matrix_02[i * 49:(i + 1) * 49, i] = 0.5
for i in range(4):
    if i != 3:
        matrix_20[196 + i * 12:196 + (i + 1) * 13, i] = 0.5   # <- possibly should be matrix_02
    else:
        matrix_20[196 + i * 12:, i] = 0.5                     # <- possibly should be matrix_02

At this point, the code is modifying matrix_20 again, but based on the logic, I think it is supposed to continue assigning values to matrix_02.