小编Dar*_*oob的帖子

我正在尝试将预训练的火炬模型转换为 ONNX，但收到以下错误：

RuntimeError: step!=1 is currently not supported

我正在预训练的着色模型上尝试此操作：https://github.com/richzhang/colorization

这是我在 Google Colab 中运行的代码：

!git clone https://github.com/richzhang/colorization.git
cd colorization/
import colorizers
model = colorizer_siggraph17 = colorizers.siggraph17(pretrained=True).eval()
input_names = [ "input" ]
output_names = [ "output" ]
dummy_input = torch.randn(1, 1, 256, 256, device='cpu')
torch.onnx.export(model, dummy_input, "test_converted_model.onnx", verbose=True,
                  input_names=input_names, output_names=output_names)

我很感激任何帮助:)

更新 1： @Proko 建议解决了 ONNX 导出问题。现在，当我尝试将 ONNX 转换为 TensorRT 时，我遇到了一个可能相关的新问题。我收到以下错误：

[TensorRT] ERROR: Network must have at least one output

这是我使用的代码：

import torch
import pycuda.driver as cuda
import pycuda.autoinit
import tensorrt …