import deepspeed only if use_deepspeed is True

feizi · Jul 10, 2023 · 82724cc · 82724cc
1 parent 31a2e15
commit 82724cc
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tortoise/models/autoregressive.py b/tortoise/models/autoregressive.py
@@ -3,7 +3,6 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-import deepspeed
 from transformers import GPT2Config, GPT2PreTrainedModel, LogitsProcessorList
 from transformers.modeling_outputs import CausalLMOutputWithCrossAttentions
 from transformers.utils.model_parallel_utils import get_device_map, assert_device_map
@@ -362,6 +361,7 @@ def post_init_gpt2_config(self, use_deepspeed=False):
             self.mel_head
         )
         if use_deepspeed:
+            import deepspeed
             self.ds_engine = deepspeed.init_inference(model=self.inference_model,  
                                                     mp_size=1,
                                                     replace_with_kernel_inject=True,