From af68fc34ab26508247f82f9c0f84053c99a398c0 Mon Sep 17 00:00:00 2001
From: Zongyuan Yang <1371618091@qq.com>
Date: Mon, 27 May 2024 15:24:16 +0800
Subject: [PATCH] Update train.py

---
 train.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/train.py b/train.py
index d70240a..eed6bb6 100644
--- a/train.py
+++ b/train.py
@@ -8,7 +8,7 @@ from pytorch_lightning import seed_everything
 from pytorch_lightning.trainer import Trainer
 from pytorch_lightning.strategies import DDPStrategy
 from pytorch_lightning.callbacks import Callback
-from pytorch_lightning.utilities import rank_zero_only
+from pytorch_lightning.utilities import rank_zero_only, rank_zero_warn
 
 from src.utils.train_util import instantiate_from_config