Fix: Respect mask_ratio=0.0 in PrithviMAE.forward

[akshanthsaik]

Replaced the truthiness check with an explicit None check so that 0.0 is treated as a valid.

[romeokienzler]

Thanks @akshanthsaik will look into it

[Copilot]

Pull request overview

This PR updates PrithviMAE.forward to treat mask_ratio=0.0 as a valid value by replacing a truthiness-based defaulting check with an explicit None check.

Changes:

• Replace mask_ratio = mask_ratio or self.mask_ratio with an explicit None fallback to preserve 0.0.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[akshanthsaik]

Added a guard in forward_loss to handle mask_sum==0, which occurs when mask_ratio=0.0, preventing division by zero and NaNs.
I think, this ensures the forward pass remains numerically stable
@romeokienzler please take a look

[akshanthsaik]

Just realised,
Instead of silently defaulting or returning zero loss, would you consider just raising an error here?

[romeokienzler]

@ahmedemam576 reviewing

[akshanthsaik]

Addressed the reviews of copilot in 7de1639

• Restored mask loss computation while handling the mask_ratio=0.0 path safely.
• Updated the mask_ratio type annotation to reflect that None is accepted.
• Removed the previous zero-loss fallback so masked loss behavior remains unchanged for valid masks.
  @romeokienzler @ahmedemam576 Could you please take a look?
  Thanks

[romeokienzler]

Thanks @akshanthsaik — reviewed. The core fix is correct: the mask_ratio or self.mask_ratio truthiness bug is properly addressed with the explicit None check, and the float | None annotation is now accurate. The ValueError placement (after resolution) is also correct — it covers both an explicit mask_ratio=0.0 and a model configured with self.mask_ratio=0.0.

One thing worth aligning on before merge:

Does raising on 0.0 actually resolve #1096? The issue asks to run the decoder without a mask ("when I'm trying to use the decoder without a mask, my input mask value of 0.0 is ignored"). After this PR, 0.0 is now correctly delivered to the logic — but then immediately rejected. So the originally-reported use case (no-mask reconstruction/inference) still can't be done. Raising is defensible since the reconstruction loss is genuinely undefined with zero masked patches, and the discussion converged here. But could you confirm with the use case in mind whether blocking it is acceptable, vs. providing a no-loss reconstruction path? forward is also the inference entry point, yet it unconditionally computes loss.

Minor: the error message ("Use a mask_ratio > 0.0 for training") is a bit training-specific given forward is used for inference too.

No blocking correctness issues — just want to make sure we're closing #1096 the way the reporter expected.

[akshanthsaik]

Thanks, @romeokienzler

I agree, for MAE models, mask_ratio=0.0 is useful in inference (full-image reconstruction/features) but not in training.
I think we can:

• Allow mask_ratio=0.0 and skip loss computation in that case.
• Only raise ValueError if we’re in training mode and mask_sum == 0 (i.e., no tokens to learn from).

Does that match what you expect, or do you prefer a different approach?