typo fix

01dffdb5 · Vijay Korthikanti · 0bc75448 · 01dffdb5
Commit 01dffdb5 authored Oct 12, 2020 by Vijay Korthikanti
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

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

No files found.
--- a/megatron/model/transformer.py
+++ b/megatron/model/transformer.py
@@ -218,7 +218,7 @@ class ParallelSelfAttention(MegatronModule):
           if checkpoint_version == 0:
               # [s, b, (3 * np * hn)] --> [s, b, (np * 3 * hn)]
               mixed_x_layer = self._transpose_last_dim(mixed_x_layer, 3, True)
-           elif checkpoint_version == 1:
+           elif checkpoint_version == 1.0:
               # [s, b, (np * hn * 3)] --> [s, b, (np * 3 * hn)]
               mixed_x_layer = self._transpose_last_dim(mixed_x_layer, 3, False)