fix: Use len() to get size of dataset

Signed-off-by: Dheeraj Peri <[email protected]>
pytorch · Jun 18, 2021 · ccc60d5 · ccc60d5
1 parent eb39f9c
commit ccc60d5
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/py/trtorch/ptq.py b/py/trtorch/ptq.py
@@ -26,7 +26,8 @@ def get_batch_size(self):
 
 
 def get_batch(self, names):
-    if self.current_batch_idx + self.batch_size > self.data_loader.dataset.data.shape[0]:
+    print("Current batch idx: ", self.current_batch_idx, " Dataset size: ", len(self.data_loader.dataset))
+    if self.current_batch_idx + self.batch_size > len(self.data_loader.dataset):
         return None
 
     batch = self.dataset_iterator.next()