bug fix for prediction

Signed-off-by: Farhan Ahmed <[email protected]>

Author: f4str

art/estimators/classification/pytorch.py

@@ -322,7 +322,7 @@ def predict(  # pylint: disable=W0221
         dataloader = DataLoader(dataset=dataset, batch_size=batch_size, shuffle=False)
 
         results_list = []
-        for x_batch in dataloader:
+        for (x_batch,) in dataloader:
             # Move inputs to device
             x_batch = x_batch.to(self._device)
 

art/estimators/regression/pytorch.py

@@ -260,7 +260,7 @@ def predict(  # pylint: disable=W0221
         dataloader = DataLoader(dataset=dataset, batch_size=batch_size, shuffle=False)
 
         results_list = []
-        for x_batch in dataloader:
+        for (x_batch,) in dataloader:
             # Move inputs to device
             x_batch = x_batch.to(self._device)