diff --git a/modelopt/torch/utils/dataset_utils.py b/modelopt/torch/utils/dataset_utils.py
index 2f8397e7a..d4cf5049d 100644
--- a/modelopt/torch/utils/dataset_utils.py
+++ b/modelopt/torch/utils/dataset_utils.py
@@ -16,7 +16,6 @@
 """Utility functions for getting samples and forward loop function for different datasets."""
 
 import copy
-import math
 from collections.abc import Callable
 from typing import TYPE_CHECKING, Any
 from warnings import warn
@@ -206,8 +205,6 @@ def get_dataset_dataloader(
     if isinstance(dataset_name, str):
         dataset_name = [dataset_name]
 
-    num_samples = [math.ceil(num_sample / batch_size) * batch_size for num_sample in num_samples]
-
     assert len(dataset_name) == len(num_samples), (
         "dataset_name and num_samples must be the same length"
     )
diff --git a/modelopt/torch/utils/speech_dataset_utils.py b/modelopt/torch/utils/speech_dataset_utils.py
index 6ca2cdad6..0d414f7ec 100644
--- a/modelopt/torch/utils/speech_dataset_utils.py
+++ b/modelopt/torch/utils/speech_dataset_utils.py
@@ -15,7 +15,6 @@
 
 """Utility functions for getting samples and forward loop function for different speech datasets."""
 
-import math
 from typing import Any
 
 import torch
@@ -101,8 +100,6 @@ def get_speech_dataset_dataloader(
     """
     assert processor is not None, "Please provide a valid processor."
 
-    num_samples = math.ceil(num_samples / batch_size) * batch_size
-
     dataset = _get_speech_dataset(dataset_name, num_samples=num_samples)
     first_sample = next(iter(dataset))
     first_text = first_sample["text"]
diff --git a/modelopt/torch/utils/vlm_dataset_utils.py b/modelopt/torch/utils/vlm_dataset_utils.py
index 4c0d7c129..1d9f59484 100644
--- a/modelopt/torch/utils/vlm_dataset_utils.py
+++ b/modelopt/torch/utils/vlm_dataset_utils.py
@@ -15,7 +15,6 @@
 
 """Utility functions for getting samples and forward loop function for different vlm datasets."""
 
-import math
 from typing import Any
 
 from torch.utils.data import DataLoader
@@ -93,8 +92,6 @@ def get_vlm_dataset_dataloader(
     """
     assert processor is not None, "Please provide a valid processor."
 
-    num_samples = math.ceil(num_samples / batch_size) * batch_size
-
     dataset = _get_vlm_dataset(dataset_name, num_samples=num_samples)
     # Apply the preprocessing function to the dataset
     processed_dataset = dataset.map(