From e1bf187afaf04d41c040d8e364d0c340ab4337f2 Mon Sep 17 00:00:00 2001 From: gbranaa4-hue Date: Sun, 6 Sep 2026 11:54:40 -0700 Subject: [PATCH] fix(networks): AffineHead follows the module's dtype, not just device self.grid is a plain attribute (torch.stack(...).to(dtype=torch.float)), never register_buffer'd. forward() re-derived only its device (self.grid.to(device=f.device)) every call, never its dtype. Casting a GlobalNet/LocalNet (both build on AffineHead) to half precision left self.grid at float32 while theta (from self.fc, whose weights did move) became float16 -- the very first half-precision forward call crashed: RuntimeError: expected scalar type Half but found Float at affine_transform's torch.einsum, which requires matching dtypes. Fix: re-derive dtype alongside device every forward call, mirroring the input f's dtype -- the same reference point self.fc(f...) already requires matching, so this is guaranteed consistent with theta's dtype by the time affine_transform(theta) - self.grid runs. Verified: half-precision GlobalNet forward no longer crashes (both 2D and via the full network), returns float16 output at the correct shape; float32 path is unaffected (byte-identical, since f.dtype was already torch.float in that case). Added test_half_precision_forward to test_globalnet.py; fails with the predicted RuntimeError on unpatched code, passes with the fix. Full existing GlobalNet/RegUNet test suites pass unchanged. Co-Authored-By: Claude Sonnet 5 --- monai/networks/nets/regunet.py | 2 +- tests/networks/nets/test_globalnet.py | 14 ++++++++++++++ 2 files changed, 15 insertions(+), 1 deletion(-) diff --git a/monai/networks/nets/regunet.py b/monai/networks/nets/regunet.py index 4d6150ea1be..c8f9829bae2 100644 --- a/monai/networks/nets/regunet.py +++ b/monai/networks/nets/regunet.py @@ -295,7 +295,7 @@ def affine_transform(self, theta: torch.Tensor): def forward(self, x: list[torch.Tensor], image_size: list[int]) -> torch.Tensor: f = x[0] - self.grid = self.grid.to(device=f.device) + self.grid = self.grid.to(device=f.device, dtype=f.dtype) theta = self.fc(f.reshape(f.shape[0], -1)) if self.save_theta: self.theta = theta.detach() diff --git a/tests/networks/nets/test_globalnet.py b/tests/networks/nets/test_globalnet.py index ecb0243a1be..fcb0218aec1 100644 --- a/tests/networks/nets/test_globalnet.py +++ b/tests/networks/nets/test_globalnet.py @@ -98,6 +98,20 @@ def test_script(self, input_param, input_shape, _): test_data = torch.randn(input_shape) test_script_save(net, test_data) + @parameterized.expand(TEST_CASES_GLOBAL_NET) + def test_half_precision_forward(self, input_param, input_shape, expected_shape): + # AffineHead.grid was a plain float32 attribute, moved by device only + # (`self.grid.to(device=f.device)`) in forward(). Casting the network to half + # precision left `self.grid` at float32 while theta became float16, and + # `affine_transform`'s einsum raised `RuntimeError: expected scalar type Half + # but found Float` on the very first half-precision forward call. + net = GlobalNet(**input_param).half() + with eval_mode(net): + img = torch.randn(input_shape, dtype=torch.float16) + result = net(img) + self.assertEqual(result.dtype, torch.float16) + self.assertEqual(result.shape, expected_shape) + if __name__ == "__main__": unittest.main()