Removal of unneeded changes so that diff is smaller

megatron/model/bert_model.py

@@ -19,7 +19,6 @@ import torch
 
 from megatron import get_args
 from megatron import mpu
-from megatron.model.language_model import Embedding
 from megatron.model.language_model import parallel_lm_logits
 from megatron.model.language_model import get_language_model
 from megatron.model.transformer import LayerNorm

megatron/model/classification.py

@@ -56,7 +56,8 @@ class Classification(MegatronModule):
             attention_mask, next(self.language_model.parameters()).dtype)
         position_ids = bert_position_ids(input_ids)
 
-        _, pooled_output = self.language_model(input_ids, position_ids,
+        _, pooled_output = self.language_model(input_ids,
+                                               position_ids,
                                                extended_attention_mask,
                                                tokentype_ids=tokentype_ids)
 

megatron/model/gpt2_model.py

@@ -21,7 +21,6 @@ from megatron import get_args
 from megatron import mpu
 from megatron.module import MegatronModule
 
-from .language_model import Embedding
 from .language_model import parallel_lm_logits
 from .language_model import get_language_model
 from .utils import init_method_normal

megatron/model/multiple_choice.py

@@ -68,7 +68,8 @@ class MultipleChoice(MegatronModule):
             attention_mask, next(self.language_model.parameters()).dtype)
         position_ids = bert_position_ids(input_ids)
 
-        _, pooled_output = self.language_model(input_ids, position_ids,
+        _, pooled_output = self.language_model(input_ids,
+                                               position_ids,
                                                extended_attention_mask,
                                                tokentype_ids=tokentype_ids)
 

megatron/model/realm_model.py

@@ -172,7 +172,8 @@ class IREncoderBertModel(MegatronModule):
         position_ids = bert_position_ids(input_ids)
 
         lm_output, pooled_output = self.language_model(
-            input_ids, position_ids,
+            input_ids,
+            position_ids,
             extended_attention_mask,
             tokentype_ids=tokentype_ids)