export onnx will add Identity operator

[zhuyuanxiang]

When I wirte a multi-head model and export the model to onnx format, I will see the identity operator in the onnx file.

This onnx file can't be read by C++ application, and met the parse error.

I use to torch.load_state_dict() the parameters, then the onnx will be good.

input_sample = torch.Tensor(1, cfg.model.net.in_chan, *cfg.data.crop_size), "pred" ckpt_path = Path(cfg.ckpt_path) checkpoint = torch.load(ckpt_path) model.load_state_dict(checkpoint['state_dict']) onnx_path = ckpt_path.parent.joinpath(ckpt_path.stem + ".onnx") model.to_onnx( onnx_path, input_sample, verbose=True, opset_version=11, training=TrainingMode.EVAL, do_constant_folding=False, operator_export_type=OperatorExportTypes.ONNX_FALLTHROUGH, )

[zhuyuanxiang]

I had tried load_from_checkpoint(), which is lightning's function, and the identity operator will be added to onnx model.

Maybe it's lightning's bug.