Solved by Gullal in post #4 
              
              
                @Ferruolo - You should be using conv2d here instead of conv3d. 
For 3x3x1024, it means that you have 1024 kernel filters of 3x3 size each. So your out_channels here is 1024, kernel_size is 3x3 and in_channels is the number of channels from the previous layer.