Add bf16-mixed and 16-mixed in module.py (#7227)

Signed-off-by: Abhishree <abhishreetm@gmail.com>
NVIDIA · Aug 15, 2023 · c561196 · c561196
1 parent 3365cde
commit c561196
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/nemo/collections/nlp/modules/common/megatron/module.py b/nemo/collections/nlp/modules/common/megatron/module.py
@@ -263,13 +263,13 @@ def __init__(self, config: ModelParallelConfig, module, precision, share_token_e
         super().__init__(config=config, share_token_embeddings=share_token_embeddings)
         self.precision = precision
 
-        if precision == 'bf16':
+        if precision in ['bf16', 'bf16-mixed']:
             self.add_module('module', module.bfloat16())
 
             def float16_converter(val):
                 return val.bfloat16()
 
-        elif int(precision) == 16:
+        elif precision in [16, '16', '16-mixed']:
             self.add_module('module', module.half())
 
             def float16_converter(val):