fix(transformer): add model eval

gnes-ai · Sep 6, 2019 · fbfa1e4746b42a75f807d5e691249048ccf4e732 · fbfa1e4
1 parent 2066beb
commit fbfa1e4746b42a75f807d5e691249048ccf4e732
Showing with 1 addition and 1 deletion.

+1 −1 gnes/encoder/text/transformer.py
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)