| 127 | |
| 128 | |
| 129 | def save_pte_program( |
| 130 | prog: ExecutorchProgramManager, model_name: str, output_dir: str = "" |
| 131 | ) -> str: |
| 132 | if model_name.endswith(".pte"): |
| 133 | filename = model_name |
| 134 | else: |
| 135 | filename = os.path.join(output_dir, f"{model_name}.pte") |
| 136 | |
| 137 | try: |
| 138 | # Write program to file. |
| 139 | with open(filename, "wb") as file: |
| 140 | prog.write_to_file(file) |
| 141 | logging.info(f"Saved exported program to {filename}") |
| 142 | # Write data to file/s. |
| 143 | prog.write_tensor_data_to_file(outdir=output_dir) |
| 144 | except Exception as e: |
| 145 | logging.error(f"Error while saving to {filename}: {e}") |
| 146 | |
| 147 | return filename |