| 15 | |
| 16 | |
| 17 | def save_pte_program( |
| 18 | prog: ExecutorchProgramManager, model_name: str, output_dir: str = "" |
| 19 | ) -> str: |
| 20 | if model_name.endswith(".pte"): |
| 21 | filename = model_name |
| 22 | visualize_file_name = f"{model_name}.json" |
| 23 | else: |
| 24 | filename = os.path.join(output_dir, f"{model_name}.pte") |
| 25 | visualize_file_name = os.path.join(output_dir, f"{model_name}.json") |
| 26 | try: |
| 27 | with open(filename, "wb") as file: |
| 28 | prog.write_to_file(file) |
| 29 | logging.info(f"Saved exported program to {filename}") |
| 30 | except Exception as e: |
| 31 | logging.error(f"Error while saving to {filename}: {e}") |
| 32 | |
| 33 | visualize_with_clusters(prog.exported_program(), visualize_file_name, False) |
| 34 | return filename |