fixed a bug

b10ca211 · mshoeybi · 8f7c0d11 · b10ca211
Commit b10ca211 authored Oct 12, 2021 by mshoeybi
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 1 deletion

megatron/model/transformer.py megatron/model/transformer.py +2 -1

No files found.
--- a/megatron/model/transformer.py
+++ b/megatron/model/transformer.py
@@ -214,8 +214,9 @@ class ParallelAttention(MegatronModule):
        # is not provided, make sure there is no potential memory left
        # from previous inference.
        else:
+            self.inference_key_memory = None
            self.inference_value_memory = None
-            self.inference_current_sequence_len = None
+

        # =====================
        # Query, Key, and Value