Adress more review comments

d836d498 · Vijay Korthikanti · 4b3519cb · d836d498
Commit d836d498 authored Jan 12, 2021 by Vijay Korthikanti
Hide whitespace changes
Inline Side-by-side

Showing with 4 additions and 4 deletions

megatron/model/language_model.py megatron/model/language_model.py +4 -4

No files found.
--- a/megatron/model/language_model.py
+++ b/megatron/model/language_model.py
@@ -406,10 +406,10 @@ class TransformerLanguageModelBase(MegatronModule):
                state_dict_[self._pooler_key] \
                    = self.pooler.state_dict_for_save_checkpoint(
                        destination, prefix, keep_vars)
-            if self.add_decoder:
-                state_dict_[self._decoder_key] \
-                    = self.decoder.state_dict_for_save_checkpoint(
-                        destination, prefix, keep_vars)
+        if self.add_decoder:
+            state_dict_[self._decoder_key] \
+                = self.decoder.state_dict_for_save_checkpoint(
+                    destination, prefix, keep_vars)

        return state_dict_