@@ -167,7 +167,6 @@ class LRWarmups:
 class Samplers:
     """Static class to hold all the supported Samplers names"""
 
-    INFINITE = "InfiniteSampler"
     REPEAT_AUG = "RepeatAugSampler"
     DISTRIBUTED = "DistributedSampler"
     SEQUENTIAL = "SequentialSampler"

@@ -1,7 +1,6 @@
-from super_gradients.training.datasets.samplers.infinite_sampler import InfiniteSampler
 from super_gradients.training.datasets.samplers.repeated_augmentation_sampler import RepeatAugSampler
 from super_gradients.common.object_names import Samplers
 from super_gradients.common.registry.registry import SAMPLERS
 
 
-__all__ = ["SAMPLERS", "Samplers", "InfiniteSampler", "RepeatAugSampler"]
+__all__ = ["SAMPLERS", "Samplers", "RepeatAugSampler"]

# Copyright (c) Megvii, Inc. and its affiliates.
# Apache 2.0 license: https://github.com/Megvii-BaseDetection/YOLOX/blob/main/LICENSE

import itertools
from typing import Optional

import torch
import torch.distributed as dist
from torch.utils.data.sampler import Sampler
from deprecate import deprecated

from super_gradients.common.object_names import Samplers
from super_gradients.common.registry.registry import register_sampler


@register_sampler(Samplers.INFINITE)
class InfiniteSampler(Sampler):
    """
    In training, we only care about the "infinite stream" of training data.
    So this sampler produces an infinite stream of indices and
    all workers cooperate to correctly shuffle the indices and sample different indices.
    The samplers in each worker effectively produces `indices[worker_id::num_workers]`
    where `indices` is an infinite stream of indices consisting of
    `shuffle(range(size)) + shuffle(range(size)) + ...` (if shuffle is True)
    or `range(size) + range(size) + ...` (if shuffle is False)
    """

    @deprecated(target=None, deprecated_in="3.0.8", remove_in="3.1.0")
    def __init__(
        self,
        dataset,
        shuffle: bool = True,
        seed: Optional[int] = 0,
        rank=0,
        world_size=1,
    ):
        """
        :param size:    Total number of data of the underlying dataset to sample from
        :param shuffle: Whether to shuffle the indices or not
        :param seed:    Initial seed of the shuffle. Must be the same across all workers.
                If None, will use a random seed shared among workers (require synchronization among all workers).
        """
        self._size = len(dataset)
        assert len(dataset) > 0
        self._shuffle = shuffle
        self._seed = int(seed)

        if dist.is_available() and dist.is_initialized():
            self._rank = dist.get_rank()
            self._world_size = dist.get_world_size()
        else:
            self._rank = rank
            self._world_size = world_size

    def __iter__(self):
        start = self._rank
        yield from itertools.islice(self._infinite_indices(), start, None, self._world_size)

    def _infinite_indices(self):
        g = torch.Generator()
        g.manual_seed(self._seed)
        while True:
            if self._shuffle:
                yield from torch.randperm(self._size, generator=g)
            else:
                yield from torch.arange(self._size)

    def __len__(self):
        return self._size // self._world_size

@@ -32,7 +32,7 @@ from super_gradients.common.factories.losses_factory import LossesFactory
 from super_gradients.common.factories.metrics_factory import MetricsFactory
 
 from super_gradients.training import utils as core_utils, models, dataloaders
-from super_gradients.training.datasets.samplers import InfiniteSampler, RepeatAugSampler
+from super_gradients.training.datasets.samplers import RepeatAugSampler
 from super_gradients.training.exceptions.sg_trainer_exceptions import UnsupportedOptimizerFormat
 from super_gradients.training.metrics.metric_utils import (
     get_metrics_titles,
@@ -160,7 +160,6 @@ class Trainer:
         self.strict_load = StrictLoad.ON
         self.load_ema_as_net = False
         self.ckpt_best_name = "ckpt_best.pth"
-        self._infinite_train_loader = False
         self._first_backward = True
 
         # METRICS
@@ -461,11 +460,7 @@ class Trainer:
             progress_bar_train_loader.set_postfix(**pbar_message_dict)
             self.phase_callback_handler.on_train_batch_end(context)
 
-            # TODO: ITERATE BY MAX ITERS
-            # FOR INFINITE SAMPLERS WE MUST BREAK WHEN REACHING LEN ITERATIONS.
-            if (self._infinite_train_loader and batch_idx == len(self.train_loader) - 1) or (
-                self.max_train_batches is not None and self.max_train_batches - 1 <= batch_idx
-            ):
+            if self.max_train_batches is not None and self.max_train_batches - 1 <= batch_idx:
                 break
 
         self.train_monitored_values = sg_trainer_utils.update_monitored_values_dict(
@@ -1022,10 +1017,10 @@ class Trainer:
                     "You are using a SequentialSampler on you training dataloader, while working on DDP. "
                     "This cancels the DDP benefits since it makes each process iterate through the entire dataset"
                 )
-            if not isinstance(train_sampler, (DistributedSampler, InfiniteSampler, RepeatAugSampler)):
+            if not isinstance(train_sampler, (DistributedSampler, RepeatAugSampler)):
                 logger.warning(
                     "The training sampler you are using might not support DDP. "
-                    "If it doesnt, please use one of the following sampler: DistributedSampler, InfiniteSampler, RepeatAugSampler"
+                    "If it doesnt, please use one of the following sampler: DistributedSampler, RepeatAugSampler"
                 )
         self.training_params = TrainingParams()
         self.training_params.override(**training_params)
@@ -1164,10 +1159,6 @@ class Trainer:
 
         self._initialize_mixed_precision(self.training_params.mixed_precision)
 
-        self._infinite_train_loader = (hasattr(self.train_loader, "sampler") and isinstance(self.train_loader.sampler, InfiniteSampler)) or (
-            hasattr(self.train_loader, "batch_sampler") and isinstance(self.train_loader.batch_sampler.sampler, InfiniteSampler)
-        )
-
         self.ckpt_best_name = self.training_params.ckpt_best_name
 
         if self.training_params.max_train_batches is not None:

@@ -165,10 +165,9 @@ class DataLoaderFactoryTest(unittest.TestCase):
 
     def test_imagenet_resnet50_kd_train_creation(self):
         # Here we need to overwrite the sampler because the RepeatAugSampler used in KD is only supported for DDP
-        dl = imagenet_resnet50_kd_train(dataloader_params={"sampler": {"InfiniteSampler": {}}})
+        dl = imagenet_resnet50_kd_train(dataloader_params={"sampler": {"RandomSampler": {}}})
         self.assertTrue(isinstance(dl, DataLoader))
         self.assertTrue(isinstance(dl.dataset, ImageNetDataset))
-        self.assertTrue(dl.sampler._shuffle)
 
     def test_imagenet_resnet50_kd_val_creation(self):
         dl = imagenet_resnet50_kd_val()