fix

Author: tastelikefeet

src/twinkle/utils/torch_utils.py

@@ -123,7 +123,7 @@ def selective_log_softmax(logits, index) -> 'torch.Tensor':
 
     try:
         from megatron.core import parallel_state as mpu
-        if mpu.get_tensor_model_parallel_world_size() > 1:
+        if mpu.get_tensor_model_parallel_world_size() >= 1:
             return _vocab_parallel_selective_log_softmax(logits, index)
     except Exception:
         pass