Possible bug in trial count assignment in train()

[Fzzz123]

Hi,

I think there might be a small typo in train():

numTrials_1 = np.size(train_data_1, 0)
numTrials_2 = np.size(train_data_1, 0)

Should the second line use train_data_2 instead?

Thanks for sharing the code.