address review comments

Signed-off-by: Abhinav Khattar <aklife97@gmail.com>

megatron/core/pipeline_parallel/schedules.py

@@ -90,7 +90,7 @@ def get_forward_backward_func():
 
     collect_non_loss_data: TODO
 
-    enable_autocat (optional, default=False): If True, runs the
+    enable_autocast (optional, default=False): If True, runs the
         forward_step_func call inside torch.autocast context
 
     """
@@ -188,7 +188,7 @@ def forward_step(forward_step_func,
     set_input_tensor = get_attr_wrapped_model(model, "set_input_tensor")
     set_input_tensor(input_tensor)
 
-    context_manager = torch.autocast("cuda") if (enable_autocast or torch.is_autocast_enabled()) else nullcontext()
+    context_manager = torch.autocast("cuda") if enable_autocast else nullcontext()
     with context_manager:
         output_tensor, loss_func = forward_step_func(data_iterator, model)