fix: Use len() to get size of dataset

Signed-off-by: Dheeraj Peri <[email protected]>

Author: peri044

Diff for: 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()