Bug Deion The code below produces the following error: <

Notably, the below code works: <div class="highlight highlight-source-python notra

Adding output_format="torch" seems to solve the

🐛 [Bug] SpecViolationError: Node.meta _to_copy_default is missing val field when using unsqueeze about tensorrt HOT 2 OPEN

airalcorn2 commented on June 3, 2024

🐛 [Bug] SpecViolationError: Node.meta _to_copy_default is missing val field when using unsqueeze

from tensorrt.

Comments (2)

airalcorn2 commented on June 3, 2024

Notably, the below code works:

import torch
import torch_tensorrt


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

    def forward(self, topk_ind):
        gather_index = topk_ind.unsqueeze(-1)
        return gather_index


def main():
    model = Model().cuda()
    model.eval()
    topk_ind = torch.randint(8400, size=(1, 300)).cuda()
    inputs = [
        torch_tensorrt.Input(topk_ind.shape, dtype=torch.int32),
    ]
    enabled_precisions = {torch.half, torch.float32}
    trt_model = torch_tensorrt.compile(
        model,
        inputs=inputs,
        enabled_precisions=enabled_precisions,
        truncate_long_and_double=True,
        min_block_size=1,
    )


if __name__ == "__main__":
    main()

which suggests it's not the unsqueeze causing problems. The original code I was testing that raised the error was:

import torch
import torch_tensorrt


class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.num_queries = 300

    def forward(self, enc_outputs_class):
        _, topk_ind = torch.topk(
            enc_outputs_class.max(-1).values, self.num_queries, dim=1
        )
        gather_index = topk_ind.unsqueeze(-1)
        return gather_index


def main():
    model = Model().cuda()
    model.eval()
    enc_outputs_class = torch.randn(1, 8400, 80).cuda()
    inputs = [
        torch_tensorrt.Input(enc_outputs_class.shape),
    ]
    enabled_precisions = {torch.half, torch.float32}
    trt_model = torch_tensorrt.compile(
        model,
        inputs=inputs,
        enabled_precisions=enabled_precisions,
        truncate_long_and_double=True,
        min_block_size=1,
    )


if __name__ == "__main__":
    main()

so it seems like it has something to do with topk.

from tensorrt.

airalcorn2 commented on June 3, 2024

Adding output_format="torchscript" seems to solve the issue:

import torch
import torch_tensorrt


class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.num_queries = 300

    def forward(self, enc_outputs_class):
        _, topk_ind = torch.topk(
            enc_outputs_class.max(-1).values, self.num_queries, dim=1
        )
        gather_index = topk_ind.unsqueeze(-1)
        return gather_index


def main():
    model = Model().cuda()
    model.eval()
    enc_outputs_class = torch.randn(1, 8400, 80).cuda()
    inputs = [
        torch_tensorrt.Input(enc_outputs_class.shape),
    ]
    enabled_precisions = {torch.half, torch.float32}
    trt_model = torch_tensorrt.compile(
        model,
        inputs=inputs,
        enabled_precisions=enabled_precisions,
        truncate_long_and_double=True,
        min_block_size=1,
        output_format="torchscript",
    )


if __name__ == "__main__":
    main()