text: # Copyright (c) Soumith Chintala 2016,
  # All rights reserved
  #
  # Copyright 2020 Huawei Technologies Co., Ltd
  #
  # Licensed under the BSD 3-Clause License (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  # https://spdx.org/licenses/BSD-3-Clause.html
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  
  import torch
  
  
  class MultiEpochsDataLoader(torch.utils.data.DataLoader):
      def __init__(self, *args, **kwargs):
          super().__init__(*args, **kwargs)
          self._DataLoader__initialized = False
          self.batch_sampler = _RepeatSampler(self.batch_sampler)
          self._DataLoader__initialized = True
          self.iterator = super().__iter__()
  
      def __len__(self):
          return len(self.batch_sampler.sampler)
  
      def __iter__(self):
          for _ in range(len(self)):
              yield next(self.iterator)
  
  
  class _RepeatSampler(object):
      """ Sampler that repeats forever.
      Args:
          sampler (Sampler)
      """
  
      def __init__(self, sampler):
          self.sampler = sampler
  
      def __iter__(self):
          while True:
              yield from iter(self.sampler)
  