Add arg checks in torch.utils.data.Sampler classes (#6249)

ssnl · ezyang · commit efc91d8c6dbf · 2018-04-04T23:07:31.000-04:00
Fixes #6168 * add arg checks in torch.utils.data.Sampler * add check for positive-ness
diff --git a/torch/utils/data/sampler.py b/torch/utils/data/sampler.py
@@ -1,4 +1,5 @@
 import torch
+from torch._six import int_classes as _int_classes
 
 
 class Sampler(object):
@@ -82,7 +83,14 @@ class WeightedRandomSampler(Sampler):
     """
 
     def __init__(self, weights, num_samples, replacement=True):
-        self.weights = torch.DoubleTensor(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 integeral "
+                             "value, but got num_samples={}".format(num_samples))
+        if not isinstance(replacement, bool):
+            raise ValueError("replacement should be a boolean value, but got"
+                             "got replacement={}".format(replacement))
+        self.weights = torch.tensor(weights, dtype=torch.double)
         self.num_samples = num_samples
         self.replacement = replacement
 
@@ -110,6 +118,17 @@ class BatchSampler(object):
     """
 
     def __init__(self, sampler, batch_size, drop_last):
+        if not isinstance(sampler, Sampler):
+            raise ValueError("sampler should be an instance of "
+                             "torch.utils.data.Sampler, but got sampler={}"
+                             .format(sampler))
+        if not isinstance(batch_size, _int_classes) or isinstance(batch_size, bool) or \
+                batch_size <= 0:
+            raise ValueError("batch_size should be a positive integeral value, "
+                             "but got batch_size={}".format(batch_size))
+        if not isinstance(drop_last, bool):
+            raise ValueError("drop_last should be a boolean value, but got"
+                             "got drop_last={}".format(drop_last))
         self.sampler = sampler
         self.batch_size = batch_size
         self.drop_last = drop_last
