Dimension problem

[JayTheFreeman]

Dear gdlg, Wishing you all the best lately.
Thank you very much for sharing the CBP codes. The time I was trying to test them, I found an interesting result. My test codes are shown as follows:
import torch
from compact_bilinear_pooling import CountSketch, CompactBilinearPooling

input_size = 3
output_size = 3
mcb = CompactBilinearPooling(input_size, input_size, output_size).cuda()
x = torch.rand(4,32,32,3).cuda()
y = torch.rand(4,32,32,3).cuda()

z = mcb(x,y)
As you had answered in another issues named "Merge"，we should make sure the channel dimension be put in the last dimension of tensor, I had follow this advise in the code. With the codes above, I expected to get the output z with dimension (4, 32, 32, 3), however, the actual output dimension of z is (4, 3, 32, 2). I have been confused about this problem for a while.
It would be very kind of you to take some time out of your busy schedule to discuss this issue with me.