Address PR comments

Signed-off-by: Marc Romeyn <marcromeyn@gmail.com>

nemo/lightning/megatron_parallel.py

@@ -237,7 +237,6 @@ def forward(
         if microbatch_outputs:
             self.callbacks.event("on_megatron_reduce_microbatches_start", **context)
 
-            # TODO: Can this lead to issues?
             if isinstance(_loss_reduction, _ModuleStepFunction):
                 _loss_reduction = _loss_reduction(self[0])
 
@@ -251,8 +250,6 @@ def forward(
             else:
                 loss_mean = torch.tensor(0.0).cuda()
 
-        # TODO: How to handle all-reduce callbacks?
-
         self.callbacks.event("on_megatron_log_step_end", **context)
         self.callbacks.event("on_megatron_step_end", **context)
 
@@ -310,8 +307,6 @@ def wrapped_forward_step_func(dataloader_iter, model):
 
             output_tensor = _forward_step(model, batch)
 
-            # TODO: handle forward_post
-
             # callback
             self._setup_module(
                 forward_callback, batch=batch, model=self, forward_module=model, tensor=output_tensor,