From 7e23747c4151fbdaaac200eff80a7190564cc74b Mon Sep 17 00:00:00 2001
From: Jake Austin <jake-austin@berkeley.edu>
Date: Sun, 17 Jul 2022 19:24:01 -0700
Subject: [PATCH] Forgot to import nn.Module

---
 pyrad/data/dataloader.py | 2 ++
 1 file changed, 2 insertions(+)

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