Error in fx.select

[tfc-0]

Hi, I have run the same code, but got this error in this line.

For input:

from compact_bilinear_pooling import CountSketch, CompactBilinearPooling

input_size = 2048
output_size = 16000
mcb = CompactBilinearPooling(input_size, input_size, output_size).cuda()
x = torch.rand(4,input_size).cuda()
y = torch.rand(4,input_size).cuda()

z = mcb(x,y)

Output error:

Traceback (most recent call last):
  File "/compact_bilinear_pooling.py", line 273, in <module>
    z = mcb(x,y)
  File "/python3.9/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "compact_bilinear_pooling.py", line 263, in forward
    return CompactBilinearPoolingFn.apply(self.sketch1.h, self.sketch1.s, self.sketch2.h, self.sketch2.s, self.output_size, x, y, self.force_cpu_scatter_add)
  File "/python3.9/site-packages/torch/autograd/function.py", line 506, in apply
    return super().apply(*args, **kwargs)  # type: ignore[misc]
  File "/compact_bilinear_pooling.py", line 154, in forward
    im_fx = fx.select(-1, 1)
IndexError: select(): index 1 out of range for tensor of size [2, 1] at dimension 1