diff --git a/yolo/lazy.py b/yolo/lazy.py
index 8bcd1a0..d6091eb 100644
--- a/yolo/lazy.py
+++ b/yolo/lazy.py
@@ -3,6 +3,7 @@
 
 import hydra
 from lightning import Trainer
+from omegaconf.dictconfig import DictConfig
 
 # FIXME: messing with sys.path is a bad idea. Factor this out.
 project_root = Path(__file__).resolve().parent.parent
@@ -14,7 +15,7 @@
 
 
 @hydra.main(config_path="config", config_name="config", version_base=None)
-def main(cfg: Config):
+def main(cfg: DictConfig):
     callbacks, loggers, save_path = setup(cfg)
 
     trainer = Trainer(