fix (#6588)

PaddlePaddle · Aug 2, 2023 · 38a2be0 · 38a2be0
1 parent caaf09d
commit 38a2be0
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/llm/llama/modeling_pp.py b/llm/llama/modeling_pp.py
@@ -18,7 +18,7 @@
 import paddle.nn as nn
 from paddle.distributed.fleet.meta_parallel import LayerDesc, PipelineLayer
 
-from paddlenlp.transformers import PretrainedModel, ScatterOp
+from paddlenlp.transformers import PretrainedModel
 from paddlenlp.transformers.llama.modeling import (
     LlamaConfig,
     LlamaDecoderLayer,
@@ -94,6 +94,8 @@ def forward(self, args):
         input_ids, attention_mask, position_ids = parse_args(args)
         input_embeds = self.embed_tokens(input_ids)
         if self.sequence_parallel:
+            from paddlenlp.transformers import ScatterOp
+
             input_embeds = paddle.transpose(input_embeds, perm=[1, 0, 2])
             input_embeds = ScatterOp.apply(input_embeds)