❓ [Question] When using torch_tensorrt.compile to optimize Mask2Former's multi_scale_deformable_attn layer, an error occurs.

❓ Question

I was preparing to export a TRT model for Mask2Former using the command optimized_model = torch_tensorrt.compile(model, inputs=imgs, enabled_precisions={torch.half}), where model is a Mask2Former loaded through mmseg.
However, I encountered an error at the line value_l_ = value_list[0].flatten(2).transpose(1, 2).reshape(4 * 8, 32, 16, 16):
The error message was:
"Failed running call_method reshape(*(FakeTensor(..., device='cuda:0', size=(1, 256, 256), grad_fn=<TransposeBackward0>), 32, 32, 16, 16), **{}): shape '[32, 32, 16, 16]' is invalid for input of size 65536"

The original code was value_l_ = value_list[level].flatten(2).transpose(1, 2).reshape(bs * num_heads, embed_dims, H_, W_). Even after fixing all variables with constants, During training, this can be reshaped normally, but the above error occurs when using torch_tensorrt.compile.

Environment

Build information about Torch-TensorRT can be found by turning on debug messages

• pytorch: 2.3.0
• torch_tensorrt: 2.3.0
• OS: ubuntu20:
• mmsegmentation: 1.2.1

Additional context

The complete code is as follows:

    value_list = value.split([16*16,32*32,64*64], dim=1)
    value_l_ = value_list[0].flatten(2).transpose(1, 2).reshape(4 * 8, 32, 16, 16)
    sampling_grid_l_ = sampling_grids[:, :, :,0].transpose(1, 2).flatten(0, 1)
    sampling_value_l_ = F.grid_sample(
            value_l_,
            sampling_grid_l_,
            mode='bilinear',
            padding_mode='zeros',
            align_corners=False)
    sampling_value_list.append(sampling_value_l_)