From 40a32f46f0bf04d45ef47a2037cebe120c77afc5 Mon Sep 17 00:00:00 2001
From: wangbinluo <2538539015@qq.com>
Date: Tue, 23 Apr 2024 13:06:27 +0000
Subject: [PATCH] fix the opt upgrade

---
 colossalai/shardformer/modeling/opt.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/colossalai/shardformer/modeling/opt.py b/colossalai/shardformer/modeling/opt.py
index dd3afecd6723..76534b5d5d2e 100644
--- a/colossalai/shardformer/modeling/opt.py
+++ b/colossalai/shardformer/modeling/opt.py
@@ -43,7 +43,7 @@ def _get_attention_mask(
             is_causal=True,
         )
     else:
-        attention_mask = self.decoder._prepare_decoder_attention_mask(
+        attention_mask = _prepare_4d_causal_attention_mask(
             attention_mask,
             (batch_size, seq_length),
             hidden_states,