There are questions about the parameters needed for forward propagation in model transformation about torch2trt HOT 2 OPEN

import torch
from torch2trt import torch2trt
from torch2trt import TRTModule

device = torch.device("cuda")
print("CUDA visible devices: " + str(torch.cuda.device_count()))
print("CUDA Device Name: " + str(torch.cuda.get_device_name(device)))

model_path = "/home/jetson/Workspace/aicam/ircamera/vit_b_slim_step2_.pth"
SlimSAM_model = torch.load(model_path)
SlimSAM_model.image_encoder = SlimSAM_model.image_encoder.module
SlimSAM_model.to(device)
SlimSAM_model.eval()
print("model_path:",model_path)

SlimSAM_model.eval().cuda()

image = torch.zeros((3, 1024, 1024)).cuda() 
original_size = (1024, 1024)  
point_coords = torch.zeros((1, 5, 2)).cuda()  
point_labels = torch.zeros((1, 5)).cuda()  
boxes = torch.zeros((1, 4)).cuda() 
mask_inputs = torch.zeros((1, 1, 1024, 1024)).cuda()  

batched_input = [{
    'image': image,
    'original_size': original_size,
    'point_coords': point_coords,
    'point_labels': point_labels,
    'boxes': boxes,
    'mask_inputs': mask_inputs
}]

multimask_output = False

model_trt = torch2trt(SlimSAM_model, [batched_input, multimask_output], fp16_mode=True)

torch.save(model_trt.state_dict(), "./vit_b_slim_step2_trt.pth")

model_trt_engine = TRTModule()
model_trt_engine.load_state_dict(model_trt.state_dict())
torch.save(model_trt_engine.state_dict(), "vit_b_slim_step2_trt.engine")

CUDA visible devices: 1
CUDA Device Name: Orin
model_path: /home/jetson/Workspace/aicam/ircamera/vit_b_slim_step2_.pth
Traceback (most recent call last):
  File "convert-sam-trt_copy3.py", line 54, in <module>
    model_trt = torch2trt(SlimSAM_model, [batched_input, multimask_output], fp16_mode=True)
  File "/home/jetson/miniconda3/envs/sam0/lib/python3.8/site-packages/torch2trt-0.5.0-py3.8.egg/torch2trt/torch2trt.py", line 558, in torch2trt
    outputs = module(*inputs)
  File "/home/jetson/miniconda3/envs/sam0/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1111, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/jetson/miniconda3/envs/sam0/lib/python3.8/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
    return func(*args, **kwargs)
  File "/home/jetson/Workspace/aicam/ircamera/segment_anything_kd/modeling/sam.py", line 98, in forward
    image_embeddings = self.image_encoder(input_images)
  File "/home/jetson/miniconda3/envs/sam0/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1111, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/jetson/Workspace/aicam/ircamera/segment_anything_kd/modeling/image_encoder.py", line 136, in forward
    qkv_emb1 = torch.cat([qkv_emb1,qkv_outputs[num]],dim=0)
RuntimeError: Sizes of tensors must match except in dimension 0. Expected size 300 but got size 396 for tensor number 1 in the list.