From f923218133a6cd69e93baac043eb15a3f66f1653 Mon Sep 17 00:00:00 2001 From: Fatwir Sheikh Mohammed <81345858+fatwir@users.noreply.github.com> Date: Fri, 22 Aug 2025 04:03:44 +0530 Subject: [PATCH] Update huggingface_mae.py Adjust Fourier loss calculation for different dimensions --- huggingface_mae.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/huggingface_mae.py b/huggingface_mae.py index 3bbefae..1f697b1 100644 --- a/huggingface_mae.py +++ b/huggingface_mae.py @@ -219,7 +219,10 @@ def compute_MAE_loss( if not self.mask_fourier_loss: floss = floss.mean() else: - floss = floss.mean(dim=-1) + if floss.dim() == 4: + floss = floss.mean(dim=(-2, -1)) + else: + floss = floss.mean(dim=-1) floss = (floss * mask).sum() / mask.sum() loss_dict[self.FOURIER_LOSS] = floss.item()