I am wondering how to use the trained model for inference. The trained model from the train branch is a regular pytorch checkpoint which is different than the onnx checkpoint used for inference in the inference branch. What are the steps to use trained checkpoints for inference?