Comment by potac

Thanks. What was confusing me is the kernel size 4. Normally in (2D) convolutions you have (in_channels, out_channels, k, k) for a kxk kernel size. In the example above it the k is the first dimension instead of the last. This is in PyTorch, not sure about Keras