Map to torchtext device parameter (#55)

* Map to torchtext device parameter Co-authored-by: [email protected] <[email protected]> Co-authored-by: Sidharth Mudgal <[email protected]>
anhaidgroup · Apr 20, 2020 · 26fd492 · 26fd492
1 parent 9123fa3
commit 26fd492
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/deepmatcher/data/iterator.py b/deepmatcher/data/iterator.py
@@ -17,13 +17,14 @@ def __init__(self,
                  train,
                  batch_size,
                  sort_in_buckets=None,
+                 device=None,
                  **kwargs):
         if sort_in_buckets is None:
             sort_in_buckets = train
         self.sort_in_buckets = sort_in_buckets
         self.train_info = train_info
         super(MatchingIterator, self).__init__(
-            dataset, batch_size, train=train, repeat=False, sort=False, **kwargs)
+            dataset, batch_size, train=train, repeat=False, sort=False, device=device, **kwargs)
 
     @classmethod
     def splits(cls, datasets, batch_sizes=None, **kwargs):