Using images with multiple channels for the MeanTeacherTraine causes an error #565
Description
When training a model that expects images with multiple (in this case 2) channels, results in an error:
RuntimeError: stack expects each tensor to be equal size, but got [2, 640, 640] at entry 0 and [1, 640, 640] at entry 2.
It originates in the logger at line 57, because it tries to stack the images with the predictions and pseudo-labels.
As a work around, changing lines 49 and 50 to
torch_em.transform.raw.normalize(x1[0, 0:1]),
torch_em.transform.raw.normalize(x2[0, 0:1]),
just as is done for the predictions and pseudo-labels, makes the code run.