fix: deactivation of model evaluation mode

Modalities · May 17, 2024 · df76a09 · df76a09
1 parent 435ab82
commit df76a09
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)