fix num_tokens in JsonlDataset (InternLM#1567)

Author: jayhenry

xtuner/v1/datasets/jsonl.py

@@ -400,6 +400,8 @@ def __init__(
         else:
             self.sampled.extend(random.sample(_sampled, _target_num_samples - len(self.sampled)))
 
+        if num_tokens is not None:
+            num_tokens = num_tokens[self.sampled]
         self.num_tokens = num_tokens
         self.offsets = offsets[self.sampled]