Adressing more review comments

megatron/checkpointing.py

@@ -165,7 +165,11 @@ def save_checkpoint(iteration, model, optimizer, lr_scheduler):
 
 
 def load_checkpoint(model, optimizer, lr_scheduler, load_arg='load', strict=True):
-    """Load a model checkpoint and return the iteration."""
+    """Load a model checkpoint and return the iteration.
+    strict (bool): whether to strictly enforce that the keys in
+        :attr:`state_dict` of the checkpoint match the names of
+        parameters and buffers in model.
+    """
     args = get_args()
     load_dir = getattr(args, load_arg)
 

megatron/model/vit_model.py

@@ -13,7 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""BERT model."""
+"""Vision Transformer(VIT) model."""
 
 import math
 import einops

tasks/vision/classification.py

@@ -13,7 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""GLUE finetuning/evaluation."""
+"""Vision-classification finetuning/evaluation."""
 
 from megatron import get_args
 from megatron import print_rank_0