1 files changed, 2 insertions, 4 deletions

diff --git a/training/functional.py b/training/functional.py
index b7ea90d..78a2b10 100644
--- a/training/functional.py
+++ b/training/functional.py
@@ -353,7 +353,6 @@ def train_loop(
     train_dataloader: DataLoader,
     val_dataloader: Optional[DataLoader],
     loss_step: LossCallable,
-    no_val: bool = False,
     sample_frequency: int = 10,
     checkpoint_frequency: int = 50,
     global_step_offset: int = 0,
@@ -472,7 +471,7 @@ def train_loop(
 
             on_after_epoch(lr_scheduler.get_last_lr()[0])
 
-            if val_dataloader is not None and not no_val:
+            if val_dataloader is not None:
                 model.eval()
 
                 cur_loss_val = AverageMeter()
@@ -616,8 +615,7 @@ def train(
         optimizer=optimizer,
         lr_scheduler=lr_scheduler,
         train_dataloader=train_dataloader,
-        val_dataloader=val_dataloader,
-        no_val=no_val,
+        val_dataloader=val_dataloader if not no_val else None,
         loss_step=loss_step_,
         sample_frequency=sample_frequency,
         checkpoint_frequency=checkpoint_frequency,