diff --git a/models/export.py b/models/export.py index 42fd577..1ff9526 100644 --- a/models/export.py +++ b/models/export.py @@ -30,7 +30,7 @@ if __name__ == '__main__': # TorchScript export try: - print('\nStarting TorchScript export with torch %s...' % torch.__version) + print('\nStarting TorchScript export with torch %s...' % torch.__version__) f = opt.weights.replace('.pt', '.torchscript') # filename ts = torch.jit.trace(model, img) ts.save(f) @@ -52,6 +52,9 @@ if __name__ == '__main__': onnx_model = onnx.load(f) # load onnx model onnx.checker.check_model(onnx_model) # check onnx model print(onnx.helper.printable_graph(onnx_model.graph)) # print a human readable model - print('ONNX export success, saved as %s\nView with https://github.com/lutzroeder/netron' % f) + print('ONNX export success, saved as %s' % f) except Exception as e: print('ONNX export failed: %s' % e) + + # Finish + print('\nExports complete. Visualize with https://github.com/lutzroeder/netron.') \ No newline at end of file