import sys
import torch
from src.body import Body


def main(pt_path, onnx_path):
    body_estimation = Body(pt_path)
    torch.onnx.export(
        body_estimation.model.cpu(), 
        torch.rand(1, 3, 224, 224), 
        onnx_path,
        opset_version=11)


if __name__ == "__main__":
    main(sys.argv[1], sys.argv[2])