From 8733d4932225b557bbf9757a793a87f7859bde64 Mon Sep 17 00:00:00 2001
From: joncrall <erotemic@gmail.com>
Date: Thu, 23 Jan 2025 14:43:41 -0500
Subject: [PATCH] fix: type error in main

---
 yolo/lazy.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

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(