Reset rotary embeddings for chained inference

OpenNMT · Mar 20, 2024 · f0d012e · f0d012e
1 parent 8479526
commit f0d012e
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/onmt/modules/multi_headed_attn.py b/onmt/modules/multi_headed_attn.py
@@ -464,6 +464,12 @@ def forward(
                     or query.dtype != torch.float16
                 ):
                     if self.max_relative_positions == -1:  # Rotary Embeddings
+                        if step == 0:
+                            self.rope, self.cos, self.sin = rotaryembeddings(
+                                self.rotary_dim,
+                                base=self.rotary_theta,
+                                device=self.rope.device,
+                            )
                         if seqlen + start_pos > self.rope.size(0):
                             # Resize rotary embeddings.
                             self.rope, self.cos, self.sin = rotaryembeddings(