I suggest you use the functional API for pooling in the forward pass so that you don’t have to redefine the layers each time. Then you are flexible to just pass the size that you need during your forward pass. I.e.,
x = torch.nn.functional.max_pool2d(x, your_kernel_size, your_stride )