(self, input_param, input_shape, expected_shape)
| 61 | @parameterized.expand(TEST_CASES_CAB) |
| 62 | @skipUnless(has_einops, "Requires einops") |
| 63 | def test_shape(self, input_param, input_shape, expected_shape): |
| 64 | net = CABlock(**input_param) |
| 65 | with eval_mode(net): |
| 66 | result = net(torch.randn(input_shape)) |
| 67 | self.assertEqual(result.shape, expected_shape) |
| 68 | |
| 69 | @skipUnless(has_einops, "Requires einops") |
| 70 | def test_invalid_spatial_dims(self): |