我正在尝试将自定义 PyTorch 模型导出到 ONNX 以执行推理，但没有成功......这里的棘手之处在于，我正在尝试使用示例中所示的基于脚本的导出器，以便从我的模型。
我可以毫无怨言地导出模型，但是在尝试启动时出现InferenceSession以下错误：
Fail: [ONNXRuntimeError] : 1 : FAIL : Load model from ner.onnx failed:Type Error: Type parameter (T) bound to different types (tensor(int64) and tensor(float) in node (Concat_1260).
我试图确定该问题的根本原因，它似乎是通过使用torch.matmul()以下函数生成的（非常讨厌，因为我试图只使用 pytorch 运算符）：
@torch.jit.script
def valid_sequence_output(sequence_output, valid_mask):
X = torch.where(valid_mask.unsqueeze(-1) == 1, sequence_output, torch.zeros_like(sequence_output))
bs, max_len, _ = X.shape
tu = torch.unique(torch.nonzero(X)[:, :2], dim=0)
batch_axis = tu[:, 0]
rows_axis = tu[:, 1]
a = torch.arange(bs).repeat(batch_axis.shape).reshape(batch_axis.shape[0], -1)
a = torch.transpose(a, 0, 1)
T = torch.cumsum(torch.where(batch_axis == a, torch.ones_like(a), torch.zeros_like(a)), dim=1) - 1
cols_axis = T[batch_axis, torch.arange(batch_axis.shape[0])]
A = torch.zeros((bs, max_len, max_len))
A[(batch_axis, cols_axis, rows_axis)] = 1.0
valid_output = torch.matmul(A, X)
valid_attention_mask = torch.where(valid_output[:, :, 0] != 0, torch.ones_like(valid_mask),
torch.zeros_like(valid_mask))
return valid_output, valid_attention_mask
似乎torch.matmul不支持（根据文档）所以我尝试了一堆解决方法（例如A.matmul(X)，torch.baddbmm）但我仍然遇到同样的问题......
有关如何解决此行为的任何建议都很棒：D 感谢您的帮助！