perf(tuner): inference mode for torch evaluation (#89)

jina-ai · Oct 5, 2021 · 68fc783 · 68fc783
1 parent 91587d8
commit 68fc783
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/finetuner/tuner/pytorch/__init__.py b/finetuner/tuner/pytorch/__init__.py
@@ -53,9 +53,10 @@ def _eval(self, data, description: str = 'Evaluating'):
 
         with ProgressBar(description, message_on_done=get_desc_str) as p:
             for inputs, label in data:
-                outputs = self.wrapped_model(*inputs)
-                loss = self.wrapped_model.loss_fn(outputs, label)
-                metric = self.wrapped_model.metric_fn(outputs, label)
+                with torch.inference_mode():
+                    outputs = self.wrapped_model(*inputs)
+                    loss = self.wrapped_model.loss_fn(outputs, label)
+                    metric = self.wrapped_model.metric_fn(outputs, label)
 
                 losses.append(loss.item())
                 metrics.append(metric.numpy())