smh

basedformer/models/ds_strats.py

@@ -3,6 +3,8 @@ from torch.nn.parameter import Parameter
 from basedformer import models
 
 def GPTJTransform(model):
+    import deepspeed
+    from deepspeed.module_inject import DSPolicy
     
     class BasedformerGPTJLayerPolicy(DSPolicy):
         _orig_layer_class = None
@@ -52,8 +54,6 @@ def GPTJTransform(model):
     model.forward = model.forward_ds
     model.get_embeds = model.get_embeds_ds
 
-    import deepspeed
-    from deepspeed.module_inject import DSPolicy
     model = deepspeed.init_inference(
                 model,
                 mp_size=1,