Open mortal-Zero opened 1 week ago
Oh yes, I can execute the following code correctly.
import torch
import torch.nn as nn
from torch2trt import torch2trt
model = nn.Sequential(
nn.Conv2d(in_channels=6, out_channels=32,
kernel_size=3, stride=1,
padding=1, bias=True),
nn.BatchNorm2d(num_features=32),
nn.LeakyReLU(),
)
model.to("cuda:0").eval()
x = torch.zeros([1, 6, 96, 96]).to("cuda:0")
y = model(x)
print("=====>> input: {} || output: {}".format(x.shape, y.shape))
model_trt = torch2trt(model, [x])
Hello, and thank you for your outstanding project. I encountered an error when converting a structure containing ConvTranspose2d using torch2trt. Here is the code and the error.
Looking forward to your reply.