test: compute_log_probs returns [seq, 1]

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Author: HJSang

tests/fast/backends/training_utils/loss_hub/test_vocab_parallel_cross_entropy.py

@@ -99,7 +99,8 @@ def _check_compute_log_probs_without_megatron(rank: int, world_size: int, port:
         group = dist.new_group([0])
         logits = _full_logits(torch.float32)[:, 0]
         tokens = _target()[:, 0]
-        expected = torch.log_softmax(logits, dim=-1).gather(-1, tokens.unsqueeze(-1)).squeeze(-1)
+        # [seq, 1], the shape compute_log_probs has always returned
+        expected = torch.log_softmax(logits, dim=-1).gather(-1, tokens.unsqueeze(-1))
 
         torch.testing.assert_close(compute_log_probs(logits, tokens, group), expected, atol=1e-6, rtol=1e-6)
     finally: