fix(transformer): add model eval

gnes-ai · Sep 6, 2019 · fbfa1e4 · fbfa1e4
1 parent 2066beb
commit fbfa1e4
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/gnes/encoder/text/transformer.py b/gnes/encoder/text/transformer.py
@@ -50,7 +50,7 @@ def post_init(self):
               (RobertaModel, RobertaTokenizer, 'roberta-base')]}[self.model_name]
 
         def load_model_tokenizer(x):
-            return model_class.from_pretrained(x), tokenizer_class.from_pretrained(x)
+            return model_class.from_pretrained(x).eval(), tokenizer_class.from_pretrained(x)
 
         try:
             self.model, self.tokenizer = load_model_tokenizer(self.work_dir)