how to process the model output in the validation phase

[mpapadomanolaki]

Hello,

Thank you very much for your code.
I have a question.
Could you explain me how can I process the model output in the validation/testing phase?
Because the ground truth shape is (batch_size, 1, height, width), however the model output is (batch_size, 64, height, width).
So when the confusion matrix is calculated in line 121 of util/metrics.py, an error is thrown.

The actual error is:
Traceback (most recent call last):
File "/home/mariapap/CODE/STANet/./train.py", line 170, in
miou_current = val(opt, model)
File "/home/mariapap/CODE/STANet/./train.py", line 86, in val
score = model.test(val=True) # run inference
File "/home/mariapap/CODE/STANet/models/CDFA_model.py", line 79, in test
metrics.update(self.L.detach().cpu().numpy(), pred.detach().cpu().numpy())
File "/home/mariapap/CODE/STANet/util/metrics.py", line 121, in update
self.confusion_matrix += self.__fast_hist(lt.flatten(), lp.flatten())
File "/home/mariapap/CODE/STANet/util/metrics.py", line 108, in __fast_hist
hist = np.bincount(self.num_classes * label_gt[mask].astype(int) + label_pred[mask],
IndexError: boolean index did not match indexed array along dimension 0; dimension is 4194304 but corresponding boolean dimension is 65536

Thank you for your time.