| 85 | class TestBasicUNET(unittest.TestCase): |
| 86 | @parameterized.expand(CASES_1D + CASES_2D + CASES_3D) |
| 87 | def test_shape(self, input_param, input_shape, expected_shape): |
| 88 | device = "cuda" if torch.cuda.is_available() else "cpu" |
| 89 | print(input_param) |
| 90 | net = BasicUNet(**input_param).to(device) |
| 91 | with eval_mode(net): |
| 92 | result = net(torch.randn(input_shape).to(device)) |
| 93 | self.assertEqual(result.shape, expected_shape) |
| 94 | |
| 95 | def test_script(self): |
| 96 | net = BasicUNet(spatial_dims=2, in_channels=1, out_channels=3) |