diff --git a/relax/backends/megatron/cp_utils.py b/relax/backends/megatron/cp_utils.py
index 98129f001..d088a3558 100644
--- a/relax/backends/megatron/cp_utils.py
+++ b/relax/backends/megatron/cp_utils.py
@@ -194,7 +194,7 @@ def get_cp_local_num_tokens(
     """
     cp_size = dynamic_cp_size if dynamic_cp_size is not None else mpu.get_context_parallel_world_size()
     if cp_size == 1:
-        return sum([torch.clamp_min(loss_mask.sum(), 1) for loss_mask in loss_masks])
+        return sum([torch.clamp_min(loss_mask.sum(), 1) for loss_mask in loss_masks]).to(torch.int)
 
     # cp_size > 1: mirror the chunk slicing done in get_sum_of_sample_mean so the
     # counted tokens exactly match the ones sum_of_token contributes on this rank.
@@ -221,8 +221,8 @@ def get_cp_local_num_tokens(
 
     if total is None:
         # No samples on this rank: mirror the empty-sum behaviour of cp_size == 1.
-        return sum([loss_mask.sum() for loss_mask in loss_masks])
-    return total
+        return sum([loss_mask.sum() for loss_mask in loss_masks]).to(torch.int)
+    return total.to(torch.int)
 
 
 def all_gather_with_cp(
diff --git a/tests/backends/megatron/test_rloo_policy_loss_dispatch.py b/tests/backends/megatron/test_rloo_policy_loss_dispatch.py
index a06d45c7b..fc1bfb962 100644
--- a/tests/backends/megatron/test_rloo_policy_loss_dispatch.py
+++ b/tests/backends/megatron/test_rloo_policy_loss_dispatch.py
@@ -167,7 +167,8 @@ def test_rloo_unequal_lengths_use_global_token_scalar_and_gradient_oracle(monkey
     final_loss = token_sum_loss / normalizer
 
     assert torch.allclose(token_sum_loss, expected_token_sum)
-    assert torch.allclose(normalizer, expected_num_tokens)
+    assert normalizer.dtype == torch.int
+    assert torch.allclose(normalizer.to(expected_num_tokens.dtype), expected_num_tokens)
     assert torch.allclose(final_loss, expected_final_loss)
 
     first_length = response_lengths[0]