From d990ddd25d335a48dae218d89c6c66b0a52a5c49 Mon Sep 17 00:00:00 2001 From: gbranaa4-hue Date: Sun, 6 Sep 2026 11:50:03 -0700 Subject: [PATCH] fix(losses): BarlowTwinsLoss preserves float16/bfloat16 input dtype c_diff = (c - torch.eye(c.size(0), device=c.device)).pow_(2) built the identity matrix with no dtype=, defaulting to the global default (float32). Subtracting it from `c` (which correctly follows the input dtype) silently upcast the entire loss computation to float32 whenever the input was float16 or bfloat16 -- float32/float64 inputs were unaffected since promotion only goes up. Add dtype=c.dtype to the torch.eye() call. The second torch.eye() call (used only to build a boolean mask via .bool()) is unaffected regardless of its dtype, since torch.eye()'s 0/1 values produce the same boolean mask at any dtype -- left as-is. Verified: float16/bfloat16/float32/float64 all now preserve their input dtype; existing float32 test cases are byte-identical (all 5 TEST_CASES plus ill_shape/ill_batch_size/with_cuda pass unchanged). Added test_preserves_input_dtype, parametrized over all four dtypes; fails on unpatched code for float16/bfloat16 (returns float32), passes with the fix. Co-Authored-By: Claude Sonnet 5 --- monai/losses/barlow_twins.py | 2 +- tests/losses/test_barlow_twins_loss.py | 11 +++++++++++ 2 files changed, 12 insertions(+), 1 deletion(-) diff --git a/monai/losses/barlow_twins.py b/monai/losses/barlow_twins.py index 699594493cb..dd38577eb35 100644 --- a/monai/losses/barlow_twins.py +++ b/monai/losses/barlow_twins.py @@ -78,7 +78,7 @@ def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor: c = torch.mm(input_norm.t(), target_norm) / batch_size # input_norm.t() is FxB, target_norm is BxF so c is FxF # loss - c_diff = (c - torch.eye(c.size(0), device=c.device)).pow_(2) # FxF + c_diff = (c - torch.eye(c.size(0), dtype=c.dtype, device=c.device)).pow_(2) # FxF c_diff[~torch.eye(c.size(0), device=c.device).bool()] *= lambd_tensor return c_diff.sum() diff --git a/tests/losses/test_barlow_twins_loss.py b/tests/losses/test_barlow_twins_loss.py index 81f4032e0c7..e41296bf188 100644 --- a/tests/losses/test_barlow_twins_loss.py +++ b/tests/losses/test_barlow_twins_loss.py @@ -104,6 +104,17 @@ def check_warning_raised(self): with self.assertWarns(Warning): BarlowTwinsLoss(lambd=5e-3, batch_size=1) + @parameterized.expand([(torch.float64,), (torch.float32,), (torch.bfloat16,), (torch.float16,)]) + def test_preserves_input_dtype(self, dtype): + # The cross-correlation matrix `c` follows the input dtype, but `c - torch.eye(...)` + # used to silently upcast to float32 whenever dtype < float32, because torch.eye() + # was never given an explicit dtype and defaults to the global default (float32). + loss = BarlowTwinsLoss(lambd=5e-3) + i = torch.randn(4, 8, dtype=dtype) + j = torch.randn(4, 8, dtype=dtype) + output = loss(i, j) + self.assertEqual(output.dtype, dtype) + if __name__ == "__main__": unittest.main()