from tensorflow.python.compiler.tensorrt import trt_convert as trt


def convert_to_trt(input_model_path: str, trt_model_path: str) -> None:
    """Utility to convert and save an input SavedModel to an optimized TensorRT graph.

    Args:
        input_model_path: Path to the SavedModel to optimize.
        trt_model_path: Path to save the converted TensorRT graph.
    """
    converter = trt.TrtGraphConverterV2(
        input_saved_model_dir=input_model_path,
        precision_mode=trt.TrtPrecisionMode.FP32,
        max_workspace_size_bytes=8000000000,
    )
    converter.convert()
    converter.save(output_saved_model_dir=trt_model_path)
    print("Done Converting to TF-TRT FP32")