diff --git a/monai/networks/nets/regunet.py b/monai/networks/nets/regunet.py index 4d6150ea1b..c8f9829bae 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 ecb0243a1b..fcb0218aec 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()