don't look up dataset on val_loader if val_loader is None

michalwols · Mar 23, 2024 · ac2ac62 · ac2ac62
1 parent b289b4c
commit ac2ac62
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/yann/train/trainer.py b/yann/train/trainer.py
@@ -8,7 +8,8 @@
 
 import torch
 import torch.nn
-from torch.cuda.amp import autocast, GradScaler
+from torch import autocast
+from torch.cuda.amp import GradScaler, autocast
 from torch.optim.optimizer import Optimizer
 from torch.utils.data import Sampler, DataLoader
 from typing_extensions import Literal
@@ -267,7 +268,8 @@ def __init__(
       callbacks: Union[
         Sequence['yann.callbacks.Callback'],
         'yann.callbacks.Callbacks',
-        None
+        None,
+        bool
       ] = None,
       device: Union[torch.device, str, None] = None,
       dtype: Optional[torch.dtype] = None,
@@ -1177,7 +1179,7 @@ def __str__(self):
 VALIDATION DATASET
 =======
 
-{self.val_loader.dataset}
+{self.val_loader.dataset if self.val_loader is not None else None}
 
 LOADER
 ======