Convert model to TensorRT fails

[janblumenkamp]

I am trying to convert the model to a torch TensorRT pre-compiled model for inference on an embedded device. I am using this script:

import torch_tensorrt
import torch

model = torch.hub.load("facebookresearch/dinov2", "dinov2_vits14")

trt_ts_module = torch_tensorrt.compile(model,
    # If the inputs to the module are plain Tensors, specify them via the `inputs` argument:
    inputs = [
        torch_tensorrt.Input( # Specify input object with shape and dtype
            min_shape=[1, 3, 224, 224],
            opt_shape=[1, 3, 224, 224],
            max_shape=[1, 3, 224, 224],
            # For static size shape=[1, 3, 224, 224]
            dtype=torch.half) # Datatype of input tensor. Allowed options torch.(float|half|int8|int32|bool)
    ],
    enabled_precisions = {torch.half}, # Run with FP16
)

torch.jit.save(trt_ts_module, "trt_torchscript_module.ts") # save the TRT embedded Torchscript

But I get the following error:

Traceback (most recent call last):
  File "[...]/torchrt_compile_dino.py", line 6, in <module>
    trt_ts_module = torch_tensorrt.compile(model,
  File "[...]/lib/python3.10/site-packages/torch_tensorrt/_compile.py", line 132, in compile
    [...]
  File "/local/scratch/jb2270/miniconda3/envs/panoptes/lib/python3.10/site-packages/torch/jit/frontend.py", line 359, in build_param_list
    raise NotSupportedError(ctx_range, _vararg_kwarg_err)
torch.jit.frontend.NotSupportedError: Compiled functions can't take variable number of arguments or use keyword-only arguments with defaults:
  File "/home/jb2270/.cache/torch/hub/facebookresearch_dinov2_main/dinov2/models/vision_transformer.py", line 290
    def forward(self, *args, is_training=False, **kwargs):
                                                 ~~~~~~~ <--- HERE
        ret = self.forward_features(*args, **kwargs)
        if is_training:

It would be great if the source files could be adapted so that a TensorRT export is possible.