diff --git a/pyrad/data/dataloader.py b/pyrad/data/dataloader.py
index 3034f1b53..286f50684 100644
--- a/pyrad/data/dataloader.py
+++ b/pyrad/data/dataloader.py
@@ -23,6 +23,8 @@ from omegaconf import ListConfig
 
 from torchtyping import TensorType
 
+import torch.nn as nn
+
 from pyrad.cameras.cameras import Camera, get_camera
 from pyrad.cameras.rays import RayBundle
 from pyrad.data.image_dataset import ImageDataset