Update test_inference.py

blue03 · Jan 31, 2023 · 2d485f0 · 2d485f0
1 parent b24058b
commit 2d485f0
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/tests/test_inference.py b/tests/test_inference.py
@@ -137,6 +137,7 @@ def main():
     if args.quantize:
         model = quantize(model, weight_bit_width=8, backend="torch")
     model.cuda()
+    torch.cuda.synchronize()
 
     with open(args.prompt_file, "r") as f:
         prompt = f.readlines()