(self, input_param, input_shape, expected_shape)
| 40 | class TestQuicknat(unittest.TestCase): |
| 41 | @parameterized.expand(TEST_CASES) |
| 42 | def test_shape(self, input_param, input_shape, expected_shape): |
| 43 | device = "cuda" if torch.cuda.is_available() else "cpu" |
| 44 | print(input_param) |
| 45 | net = Quicknat(**input_param).to(device) |
| 46 | with eval_mode(net): |
| 47 | result = net(torch.randn(input_shape).to(device)) |
| 48 | self.assertEqual(result.shape, expected_shape) |
| 49 | |
| 50 | def test_script(self): |
| 51 | net = Quicknat(num_classes=1, num_channels=1) |