add set_seed for spr path (#1242)

intel · Jul 28, 2023 · 4c75efe · 4c75efe
1 parent ca1e26e
commit 4c75efe
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/workflows/chatbot/inference/generate.py b/workflows/chatbot/inference/generate.py
@@ -817,6 +817,9 @@ def main():
         from optimum.habana.utils import set_seed
 
         set_seed(args.seed)
+    else:
+        from transformers import set_seed
+        set_seed(args.seed)
 
     tokenizer_path = (
         args.tokenizer_name if args.tokenizer_name is not None else base_model_path