torch2trt
torch2trt copied to clipboard
when use own defined convertor, why int tensor input has no _trt attribute?
here is a simple reproduce code
import os
import time
import torch
import torchvision
from torch2trt.torch2trt import *
import numpy as np
class Sparse(torch.nn.Module):
def __init__(self, embedding_size):
super().__init__()
self._embedding_size = embedding_size
self._output = torch.zeros((4, self._embedding_size))
def forward(self, x):
x = x.float()
return self._output
class Model(torch.nn.Module):
def __init__(self):
super().__init__()
self.sparse = Sparse(100)
self.linear = torch.nn.Linear(100, 200)
def forward(self, x):
y = self.sparse(x)
return self.linear(y)
@tensorrt_converter(Sparse.forward)
def convert_sparse(ctx):
module = get_arg(ctx, 'self', pos=0, default=None)
x = get_arg(ctx, 'x', pos=1, default=None)
print(x._trt) # verfiy _trt attribute
if __name__ == "__main__":
model = Sparse(100)
model.eval()
x = torch.ones((1, 3, 224, 224), dtype=torch.int32).cuda()
#x = torch.ones((1, 3, 224, 224), dtype=torch.float32).cuda() It works
model_trt = torch2trt(model, [x])
AttributeError: 'Tensor' object has no attribute '_trt'