black

OpenNMT · Dec 26, 2023 · 17373ca · 17373ca
1 parent d0ec7a8
commit 17373ca
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/onmt/modules/multi_headed_attn.py b/onmt/modules/multi_headed_attn.py
@@ -509,8 +509,12 @@ def forward(
                             dim=-2,
                         )
                     if sliding_window > 0 and key.size(2) > sliding_window:
-                        self.layer_cache[1]["keys"] = self.layer_cache[1]["keys"][:, :, 1:, :]
-                        self.layer_cache[1]["values"] = self.layer_cache[1]["values"][:, :, 1:, :]
+                        self.layer_cache[1]["keys"] = self.layer_cache[1]["keys"][
+                            :, :, 1:, :
+                        ]
+                        self.layer_cache[1]["values"] = self.layer_cache[1]["values"][
+                            :, :, 1:, :
+                        ]
                     context = self.flash_attn_with_kvcache(
                         query.transpose(1, 2),
                         self.layer_cache[1]["keys"].transpose(1, 2),