See https://pytorch.org/docs/stable/generated/torch.nn.AdaptiveAvgPool2d.html

This takes fixed output size and determines kernel size based on input. This should be possible with const generics & using avg pool from #157