add generator

PaddlePaddle · Oct 19, 2023 · bf121ba · bf121ba
1 parent 1f71c36
commit bf121ba
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/python/paddle/io/dataloader/sampler.py b/python/paddle/io/dataloader/sampler.py
@@ -349,6 +349,7 @@ class SubsetRandomSampler(Sampler):
 
     Args:
         indices (sequence): a sequence of indices
+        generator(Generator, optional): specify a generator to sample the :code:`data_source`. Default None, disabled.
 
     Examples:
 
@@ -368,15 +369,15 @@ class SubsetRandomSampler(Sampler):
             5
             1
 
-    see `paddle.io.Sampler`
     """
 
-    def __init__(self, indices):
+    def __init__(self, indices, generator=None):
         if len(indices) == 0:
             raise ValueError(
                 "The length of `indices` in SubsetRandomSampler should be greater than 0."
             )
         self.indices = indices
+        assert generator is None
 
     def __iter__(self):
         for i in randperm(len(self.indices)):