fix(model): set model to evaluation mode (#218)

hyperonym · Jun 25, 2023 · de49385 · de49385
1 parent 0288db4
commit de49385
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/basaran/model.py b/basaran/model.py
@@ -78,7 +78,7 @@ def __call__(
                 temperature=temperature,
                 top_p=top_p,
             ),
-            **kwargs
+            **kwargs,
         }
 
         # Generate completion tokens.
@@ -347,4 +347,4 @@ def load_model(
     if not model.can_generate():
         raise TypeError(f"{name_or_path} is not a text generation model")
 
-    return StreamModel(model, tokenizer)
+    return StreamModel(model.eval(), tokenizer)