fix bug when using gradient_checkpointing

modeling_telechat.py

@@ -43,8 +43,6 @@ except ImportError:
     try:
         from flash_attn.flash_attn_interface import flash_attn_varlen_func as flash_attn_unpadded_func # flashattn2
         print("# FLASH ATTENTION 2 DETECTED #")
-    r
-    r
     except ImportError:
         print("# NO FLASH ATTENTION DETECTED #")
         flash_attn_unpadded_func = None
@@ -857,6 +855,8 @@ class TELECHATTransformer(TELECHATPretrainedModel):
             if output_hidden_states:
                 all_hidden_states = all_hidden_states + (hidden_states,)
 
+            rotary_embedding=self.wpe if self.relative_encoding == 'rotary' else None
+
             if self.gradient_checkpointing and self.training:
 
                 if use_cache:
@@ -880,6 +880,7 @@ class TELECHATTransformer(TELECHATPretrainedModel):
                     head_mask[i],
                     encoder_hidden_states,
                     encoder_attention_mask,
+                    rotary_embedding
                 )
             else:
                 outputs = block(
@@ -889,7 +890,7 @@ class TELECHATTransformer(TELECHATPretrainedModel):
                     head_mask=head_mask[i],
                     encoder_hidden_states=encoder_hidden_states,
                     encoder_attention_mask=encoder_attention_mask,
-                    rotary_embedding=self.wpe if self.relative_encoding == 'rotary' else None,
+                    rotary_embedding=rotary_embedding,
                     use_cache=use_cache,
                     output_attentions=output_attentions
                 )