Simplify balanced_sampler

torchnlp/_third_party/weighted_random_sampler.py

@@ -6,8 +6,10 @@
 
 class WeightedRandomSampler(Sampler):
 
-    def __init__(self, weights, num_samples, replacement=True):
-        # NOTE: Adapted `WeightedRandomSampler` to accept `num_samples=0`.
+    def __init__(self, weights, num_samples=None, replacement=True):
+        # NOTE: Adapted `WeightedRandomSampler` to accept `num_samples=0` and `num_samples=None`.
+        if num_samples is None:
+            num_samples = len(weights)
         if not isinstance(num_samples, _int_classes) or isinstance(num_samples, bool) or \
                 num_samples < 0:
             raise ValueError("num_samples should be a positive integer "

torchnlp/samplers/balanced_sampler.py

@@ -29,6 +29,4 @@ def __init__(self, data_source, get_class=identity, get_weight=lambda x: 1, **kw
             k: sum([w for c, w in zip(classified, weighted) if k == c]) for k in set(classified)
         }
         weights = [w / class_totals[c] if w > 0 else 0.0 for c, w in zip(classified, weighted)]
-        if 'num_samples' not in kwargs:
-            kwargs['num_samples'] = len(data_source)
         super().__init__(weights=weights, **kwargs)