enforce irpa extension on save

iree-org · May 2, 2024 · f1beea9 · f1beea9
1 parent 31d4378
commit f1beea9
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/shark_turbine/aot/params.py b/shark_turbine/aot/params.py
@@ -258,7 +258,10 @@ def __init__(self):
 
     def save(self, file_path: Union[str, Path]):
         """Saves the archive."""
-        self._index.create_archive_file(str(file_path))
+        str_file_path = str(file_path)
+        if not str_file_path.endswith(".irpa"):
+            file_path = str_file_path + ".irpa"
+        self._index.create_archive_file(str_file_path)
 
     def add_tensor(self, name: str, tensor: torch.Tensor):
         """Adds an named tensor to the archive."""