add torch_dtype_from_mcore_config

Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com>

Author: akoumpa

nemo/collections/llm/gpt/model/base.py

@@ -126,6 +126,14 @@ def default_layer_spec(config: "GPTConfig") -> ModuleSpec:
         return local_layer_spec(config)
 
 
+def torch_dtype_from_mcore_config(config: TransformerConfig):
+    if config.fp16:
+        return torch.float16
+    elif config.bf16:
+        return torch.bfloat16
+    else:
+        return torch.float
+
 @dataclass
 class GPTConfig(TransformerConfig, io.IOMixin):
     # From megatron.core.models.gpt.gpt_model.GPTModel