model load

export.py

@@ -67,7 +67,7 @@ def main(checkpoint_path, device, quantize = "{ }", args = None):
     print("Loading model ...")
     t0 = time.time()
     model = _load_model(
-        checkpoint_path, device=device, precision=precision, use_tp=False)
+        checkpoint_path, args.checkpoint_dir, device=device, precision=precision, use_tp=False)
 
     device_sync(device=device)  # MKG
     print(f"Time to load model: {time.time() - t0:.02f} seconds")
@@ -149,6 +149,12 @@ def cli():
         default="not_specified",
         help="Model checkpoint path.",
     )
+    parser.add_argument(
+        "--checkpoint-dir",
+        type=Path,
+        default=None,
+        help="Model checkpoint directory.",
+    )
     parser.add_argument(
         "--output-pte-path",
         type=str,