Add new get_batch() method to ChunkDataset API (#21797)

Thiago Crepaldi · facebook-github-bot · commit 07fea3f5b67e · 2019-06-14T13:39:54.000-07:00
Summary: We plan on generating python bindings for C++ ChunkDataset API using the current Pytorch Dataloader class, which must call get_batch() instead of get_batch(size) This changes doesnt break the current API, just add one more method that will make future extensions easier (WIP) Pull Request resolved: #21797 Differential Revision: D15830522 Pulled By: soumith fbshipit-source-id: 7208f305b48bf65d2783eaff43ff57a05e62c255
diff --git a/torch/csrc/api/include/torch/data/datasets/chunk.h b/torch/csrc/api/include/torch/data/datasets/chunk.h
@@ -336,6 +336,11 @@ class ChunkDataset final
     return batch_buffer_->get_batch();
   }
 
+  /// Helper method around get_batch as `batch_size` is not strictly necessary
+  BatchType get_batch() {
+    return get_batch(options_.batch_size_);
+  }
+
   /// This will clear any internal state and starts the internal prefetching
   /// mechanism for the chunk dataset.
   void reset() override {
