update

Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>

Author: ananthsub

nemo_rl/models/megatron/train.py

@@ -174,7 +174,7 @@ def forward_with_post_processing_fn(
 
     # Apply temperature scaling only for sampling-oriented post-processors.
     # Loss computation should use unscaled logits.
-    if isinstance(post_processing_fn, (LogprobsPostProcessor, TopkLogitsPostProcessor)):
+    if isinstance(post_processing_fn, (LossPostProcessor, LogprobsPostProcessor, TopkLogitsPostProcessor)):
         apply_temperature_scaling(output_tensor, cfg)
 
     # Use type checking to dispatch to the correct post-processing method