Merge pull request #54 from pytorch-labs/simplify

rm Makefile, update fairseq2 import algorithm
pytorch · Apr 5, 2024 · ea4d370 · ea4d370
2 parents 79b7bca + 10eb8b2
commit ea4d370
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/generate.py b/generate.py
@@ -393,7 +393,11 @@ def main(
 
         # Add new CLI arg
         if quantize:
+            device_sync(device=device)
+            t0q = time.time()
             quantize_model(model, quantize)
+            device_sync(device=device)  # MKG
+            print(f"Time to quantize model: {time.time() - t0q:.02f} seconds")
 
     if is_speculative:
         draft_model = _load_model(draft_checkpoint_path, device, precision, use_tp)