Account for object arrays

Signed-off-by: Beat Buesser <[email protected]>

Author: Beat Buesser

art/estimators/object_tracking/pytorch_goturn.py

@@ -665,11 +665,12 @@ def predict(self, x: np.ndarray, batch_size: int = 128, **kwargs) -> List[Dict[s
 
         for i in range(x.shape[0]):
             if isinstance(x, np.ndarray):
-                x_i = torch.from_numpy(x[[i]]).to(self.device)
+                x_i = torch.from_numpy(x[i]).to(self.device)
             else:
-                x_i = x[[i]].to(self.device)
+                x_i = x[i].to(self.device)
 
             # Apply preprocessing
+            x_i = torch.unsqueeze(x_i, dim=0)
             x_i, _ = self._apply_preprocessing(x_i, y=None, fit=False, no_grad=False)
             x_i = torch.squeeze(x_i)