Open mortal-Zero opened 2 months 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])
Any fix please? I am having same problem
I guess the issue is relevant here:
https://github.com/NVIDIA-AI-IOT/torch2trt/blob/4e820ae31b4e35d59685935223b05b2e11d47b03/torch2trt/converters/native_converters.py#L439
Change out_channels = int(weight.shape[0])
to out_channels = int(weight.shape[1])
and reinstall the package solved the issue.
Output channel is weight.shape[1]
for ConvTransposexd
, and weight.shape[0]
for Convxd
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.