Open Thrsu opened 10 months ago
I encountered the same problem with Softmax
operator when I execute the script:
import torch
from torch.nn import Module
from torch2trt import torch2trt
input_data = torch.randn([16, 16, 0], dtype=torch.float32).cuda()
class softmax(Module):
def forward(self, *args):
return torch.nn.functional.softmax(args[0], )
model = softmax().float().eval().cuda()
model_trt = torch2trt(model, [input_data])
Description:
I encountered an UnboundLocalError when trying to convert a
MultiLabelSoftMarginLoss
module using torch2trt. The error occurs when the dim parameter is not provided to the torch.mean function within the torch2trt conversion process.Reproduce:
Here is a minimal script to reproduce the issue:
Environment