From 2d485f0610a1409f561a607e807ceb862b6d241e Mon Sep 17 00:00:00 2001
From: Xiaoyu Zhang <35585791+BBuf@users.noreply.github.com>
Date: Tue, 31 Jan 2023 13:51:46 +0800
Subject: [PATCH] Update test_inference.py

---
 tests/test_inference.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tests/test_inference.py b/tests/test_inference.py
index 95b6ef5..be6ed40 100644
--- 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()