fix: Transfer calibration data to gpu when it is not a batch

Signed-off-by: Dheeraj Peri <[email protected]>
pytorch · Jun 29, 2021 · 23739cb · 23739cb
1 parent a032c3a
commit 23739cb
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/py/trtorch/ptq.py b/py/trtorch/ptq.py
@@ -34,6 +34,8 @@ def get_batch(self, names):
     # Treat the first element as input and others as targets.
     if isinstance(batch, list):
         batch = batch[0].to(self.device)
+    else:
+        batch = batch.to(self.device)
     return [batch.data_ptr()]