def __init__(self,
in_channels,
out_channels,
dropout_p,
spatial_dims: int = 2):
super().__init__()
self.max_avg_pool = MaxAvgPool(spatial_dims=spatial_dims,
kernel_size=2)
self.conv = ConvBNActBlock(2 * in_channels,
out_channels,
dropout_p,
spatial_dims=spatial_dims)
def test_shape(self, input_param, input_shape, expected_shape):
net = MaxAvgPool(**input_param)
with eval_mode(net):
result = net(torch.randn(input_shape))
self.assertEqual(result.shape, expected_shape)
def test_shape(self, input_param, input_data, expected_shape):
net = MaxAvgPool(**input_param)
net.eval()
with torch.no_grad():
result = net(input_data)
self.assertEqual(result.shape, expected_shape)