incompatible function arguments | torch.bmm()

[msaadsaeed]

Hi! I am trying to optimize a custom model. Before, I was facing not implemented error. Trying the solution provided here. However, now I am facing this problem:

TypeError: add_matrix_multiply(): incompatible function arguments. The following argument types are supported:
    1. (self: tensorrt.tensorrt.INetworkDefinition, input0: tensorrt.tensorrt.ITensor, op0: tensorrt.tensorrt.MatrixOperation, input1: tensorrt.tensorrt.ITensor, op1: tensorrt.tensorrt.MatrixOperation) -> tensorrt.tensorrt.IMatrixMultiplyLayer

I also printed the data types:

input_a_trt, input_b_trt = trt_(ctx.network, input_a, input_b)
print(type(input_a_trt))
print(type(input_b_trt))

Output:

<class 'tensorrt.tensorrt.ITensor'>
<class 'tensorrt.tensorrt.ITensor'>

[wr0112358]

Try replacing the trt_ call with:

input_a_trt, input_b_trt = add_missing_trt_tensors(ctx.network, [input_a, input_b]) input_a_trt, input_b_trt = broadcast_trt_tensors(ctx.network, [input_a_trt, input_b_trt], len(output.shape) - 1)

[msaadsaeed]

I have replaced the code as per your recommendation:

input_a = ctx.method_args[0]
input_b = ctx.method_args[1]
input_a_trt, input_b_trt = add_missing_trt_tensors(ctx.network, [input_a, input_b])
input_a_trt, input_b_trt = broadcast_trt_tensors(ctx.network, [input_a_trt, input_b_trt], len(output.shape)-1)

Data type of input_a_trt and input_b_trt:

<class 'tensorrt.tensorrt.ITensor'>
<class 'tensorrt.tensorrt.ITensor'>

Still same error:

TypeError: add_matrix_multiply(): incompatible function arguments. The following argument types are supported:
    1. (self: tensorrt.tensorrt.INetworkDefinition, input0: tensorrt.tensorrt.ITensor, op0: tensorrt.tensorrt.MatrixOperation, input1: tensorrt.tensorrt.ITensor, op1: tensorrt.tensorrt.MatrixOperation) -> tensorrt.tensorrt.IMatrixMultiplyLayer

Invoked with: <tensorrt.tensorrt.INetworkDefinition object at 0x000001CD12822AB0>, <tensorrt.tensorrt.ITensor object at 0x000001CD2668C4F0>, <tensorrt.tensorrt.ITensor object at 0x000001CD2668C8B0>

Here is my code:

from torch2trt.torch2trt import *
from torch2trt.module_test import add_module_test

@tensorrt_converter('torch.bmm')
def convert_bmm(ctx):
    input_a = ctx.method_args[0]
    input_b = ctx.method_args[1]
    output = ctx.method_return
    # input_a_trt = trt_(ctx.network, [input_a])[0]
    # input_b_trt = trt_(ctx.network, [input_b])[0]
    input_a_trt, input_b_trt = add_missing_trt_tensors(ctx.network, [input_a, input_b])
    input_a_trt, input_b_trt = broadcast_trt_tensors(ctx.network, [input_a_trt, input_b_trt], len(output.shape)-1)
    # input_a_trt, input_b_trt = trt_(ctx.network, input_a, input_b)

    layer = ctx.network.add_matrix_multiply(input_a_trt ,input_b_trt)
    output._trt = layer.get_output(0)

class Bmm(torch.nn.Module):
    def __init__(self):
        super(Bmm, self).__init__()

    def forward(self, input_a, input_b):
        return torch.bmm(input_a, input_b)

[wr0112358]

call to add_matrix_multiply is wrong too: layer = ctx.network.add_matrix_multiply(input_a_trt, trt.MatrixOperation.NONE, input_b_trt, trt.MatrixOperation.NONE)

[msaadsaeed]

Thanks for pointing out the mistake! However now I am facing a new error: [ltWrapper.cpp::nvinfer1::rt::CublasLtWrapper::setupHeuristic::327] Error Code 2: Internal Error (Assertion cublasStatus == CUBLAS_STATUS_SUCCESS failed. )

Here is a thread related to CUBLAS error. I have all the requirements installed. However, it is also not working.

My Environment TensorRT Version: 8.2.1.8 OS: Windows CUDA: 10.2 (with both patches installed) cudnn: 8.2.1.32