Merge pull request #131 from Modalities/fix/deactivate-model-eval

fix: deactivation of model evaluation mode
Modalities · May 21, 2024 · 2b562b1 · 2b562b1
2 parents 435ab82 + df76a09
commit 2b562b1
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/src/modalities/evaluator.py b/src/modalities/evaluator.py
@@ -104,6 +104,9 @@ def evaluate(
                 evaluation_result=evaluation_result,
             )
             result_dict[data_loader.dataloader_tag] = evaluation_result
+
+        model.train()
+
         return result_dict
 
     @staticmethod

diff --git a/src/modalities/trainer.py b/src/modalities/trainer.py
@@ -181,7 +181,6 @@ def train(
                 )
                 thoughput_aggregator.remove_keys()
 
-                model.train()
                 cumulated_losses = self._reset_tracked_losses()
 
             evaluation_callback(train_step_id=train_step_id)