Getting shape mismatch when trying to run experiments

[lupinthief]

Hi. Thanks so much for the great work. This is a really exciting project.
I'm trying to reproduce the experiments but get shape mismatch errors between prediction and target when trying to calculate IOU in the training loop for both exp=BraDD_AnySat_LP and exp=BraDD_AnySat_FT. These appeared after I changed the batch size to 4 due to OOM, but persist having changed it back to 8 on a bigger node. Do you have any thoughts on how to fix this?

  File "***/conda-envs/anysat/lib/python3.9/site-packages/torchmetrics/functional/segmentation/mean_iou.py", line 77, in _mean_iou_update
    _check_same_shape(preds, target)
  File "***/conda-envs/anysat/lib/python3.9/site-packages/torchmetrics/utilities/checks.py", line 39, in _check_same_shape
    raise RuntimeError(
RuntimeError: Predictions and targets are expected to have the same shape, but got torch.Size([18432, 2]) and torch.Size([384, 48, 2]).