fix: convert DTensor output to regular tensor after loss_parallel

Author: gitlost-murali

src/forge/actors/reference_model.py

@@ -194,6 +194,8 @@ async def forward(
             if parallel_dims.tp_enabled and isinstance(logits, DTensor):
                 with loss_parallel():
                     logprobs = compute_logprobs(logits, response_tokens, align=True)
+
+                logprobs = logprobs.to_local()
             else:
                 logprobs = compute_logprobs(logits, response_tokens)
             t.step("compute_logprobs")