Correct Indices calculation compares first element with ground truth for ENTIRE batch

[apurvakokate]

Starting with line 389 in train.py:
preds = torch.cat(preds) concatenates all predictions into a 1D tensor, the predictions for training samples is returned by the method

then at line 112 preds = output_train the output is assigned to preds

then at line 207 pred=preds[idx] the prediction is indexed by the index of the batch, however as a side effect of the concatenation this returns a single value

finally,
at line 249 cor_idx = np.where(pred.cpu() == gtruth)[0] the code returns a list of indices where the single value of pred is found in the ground truth. Numpy's where function does not throw an error as it internally manages its operations.

Logically, pred at line 249 should be a list of predictions with size equal to batch size and not a single value so the correct indices can be calculated properly