RuntimeError: output with shape [32, 2048, 2048] doesn't match the broadcast shape [1, 32, 2048, 2048]

Hi, Thank you for sharing your great work!

I’m running your pre-trained model to reproduce the results but I see the runtime error below. The error seems about an input shape issue. Do you have any suggestions for this?

This is the command that I put.

python main.py --dataset_name sunrgbd --nqueries 128 --test_ckpt weights/sunrgbd_masked_ep1080.pth --test_only --enc_type masked

This is the error that I have.

Traceback (most recent call last):
  File "main.py", line 427, in <module>
    launch_distributed(args)
  File "main.py", line 415, in launch_distributed
    main(local_rank=0, args=args)
  File "main.py", line 388, in main
    test_model(args, model, model_no_ddp, criterion, dataset_config, dataloaders)
  File "main.py", line 316, in test_model
    curr_iter,
  File "/home/user/anaconda3/envs/3detr/lib/python3.6/site-packages/torch/autograd/grad_mode.py", line 49, in decorate_no_grad
    return func(*args, **kwargs)
  File "/home/user/Desktop/3detr/engine.py", line 179, in evaluate
    outputs = model(inputs)
  File "/home/user/anaconda3/envs/3detr/lib/python3.6/site-packages/torch/nn/modules/module.py", line 532, in __call__
    result = self.forward(*input, **kwargs)
  File "/home/user/Desktop/3detr/models/model_3detr.py", line 306, in forward
    enc_xyz, enc_features, enc_inds = self.run_encoder(point_clouds)
  File "/home/user/Desktop/3detr/models/model_3detr.py", line 200, in run_encoder
    pre_enc_features, xyz=pre_enc_xyz
  File "/home/user/anaconda3/envs/3detr/lib/python3.6/site-packages/torch/nn/modules/module.py", line 532, in __call__
    result = self.forward(*input, **kwargs)
  File "/home/user/Desktop/3detr/models/transformer.py", line 190, in forward
    output = layer(output, src_mask=mask, src_key_padding_mask=src_key_padding_mask, pos=pos)
  File "/home/user/anaconda3/envs/3detr/lib/python3.6/site-packages/torch/nn/modules/module.py", line 532, in __call__
    result = self.forward(*input, **kwargs)
  File "/home/user/Desktop/3detr/models/transformer.py", line 288, in forward
    return self.forward_pre(src, src_mask, src_key_padding_mask, pos, return_attn_weights)
  File "/home/user/Desktop/3detr/models/transformer.py", line 272, in forward_pre
    key_padding_mask=src_key_padding_mask)
  File "/home/user/anaconda3/envs/3detr/lib/python3.6/site-packages/torch/nn/modules/module.py", line 532, in __call__
    result = self.forward(*input, **kwargs)
  File "/home/user/anaconda3/envs/3detr/lib/python3.6/site-packages/torch/nn/modules/activation.py", line 819, in forward
    attn_mask=attn_mask)
  File "/home/user/anaconda3/envs/3detr/lib/python3.6/site-packages/torch/nn/functional.py", line 3362, in multi_head_attention_forward
    attn_output_weights += attn_mask
RuntimeError: output with shape [32, 2048, 2048] doesn't match the broadcast shape [1, 32, 2048, 2048]

[rohitgirdhar]

Which CUDA version are you using with pytorch 1.9? Can you try with the CUDA 10.2 version?

[bhkim94]

I tried with pytorch 1.9 and CUDA 10.2 and that works well. The CUDA version was the reason.

Thanks for helping! @imisra @rohitgirdhar