Looks like moving the model.train() and model.eval() here should work.

src/fibad/models/example_autoencoder.py

@@ -119,9 +119,6 @@ def train_step(self, batch):
         Current loss value
             The loss value for the current batch.
         """
-        # Set the model to train mode
-        self.train()
-
         # When we run on a supervised dataset like CIFAR10, drop the labels given by the data loader
         x = batch[0] if isinstance(batch, tuple) else batch
 

src/fibad/models/example_cnn_classifier.py

@@ -55,9 +55,6 @@ def train_step(self, batch):
         Current loss value
             The loss value for the current batch.
         """
-        # Set the model to train mode
-        self.train()
-
         inputs, labels = batch
 
         self.optimizer.zero_grad()

src/fibad/pytorch_ignite.py

@@ -138,6 +138,7 @@ def create_evaluator(model: torch.nn.Module, save_function: Callable[[torch.Tens
         Engine object which when run will evaluate the model.
     """
     device = idist.device()
+    model.eval()
     model = idist.auto_model(model)
     evaluator = create_engine("forward", device, model)
 
@@ -178,6 +179,7 @@ def create_trainer(model: torch.nn.Module, config: ConfigDict, results_directory
         Engine object that will be used to train the model.
     """
     device = idist.device()
+    model.train()
     model = idist.auto_model(model)
     trainer = create_engine("train_step", device, model)