update

kplug/modeling_kplug_s2s_patch.py

@@ -24,14 +24,23 @@ class KPlugLearnedPositionalEmbedding(nn.Embedding):
         self.offset = 2
         super().__init__(num_embeddings + self.offset, embedding_dim)
 
-    def forward(self, input_ids: torch.Tensor, past_key_values_length: int = 0):
-        """`input_ids' shape is expected to be [bsz x seqlen]."""
-
-        bsz, seq_len = input_ids.shape[:2]
+    ### 4.21.1   之后的版本 
+    # def forward(self, input_ids: torch.Tensor, past_key_values_length: int = 0):
+    #     """`input_ids' shape is expected to be [bsz x seqlen]."""
+    #
+    #     bsz, seq_len = input_ids.shape[:2]
+    #     positions = torch.arange(
+    #         past_key_values_length, past_key_values_length + seq_len, dtype=torch.long, device=self.weight.device
+    #     ).expand(bsz, -1)
+    #
+    #     return super().forward(positions + self.offset)
+
+    def forward(self, input_ids_shape: torch.Size, past_key_values_length: int = 0):
+        """`input_ids_shape` is expected to be [bsz x seqlen]."""
+        bsz, seq_len = input_ids_shape[:2]
         positions = torch.arange(
             past_key_values_length, past_key_values_length + seq_len, dtype=torch.long, device=self.weight.device
-        ).expand(bsz, -1)
-
+        )
         return super().forward(positions + self.offset)