From 78fdc372a1efd6795d49047967a78fbab8e9fb22 Mon Sep 17 00:00:00 2001
From: Richard Erkhov <104026872+RichardErkhov@users.noreply.github.com>
Date: Mon, 22 May 2023 16:22:37 +0300
Subject: [PATCH] Update train.py

---
 RWKV-v4neo/train.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/RWKV-v4neo/train.py b/RWKV-v4neo/train.py
index a2aa7e8d..d269290e 100644
--- a/RWKV-v4neo/train.py
+++ b/RWKV-v4neo/train.py
@@ -3,6 +3,7 @@
 ########################################################################################################
 
 if __name__ == "__main__":
+    import deepspeed
     from argparse import ArgumentParser
     from pytorch_lightning import Trainer
     from pytorch_lightning.utilities import rank_zero_info, rank_zero_only