[fix] accepting batch size is greater than 1, needed during training

Author: chyomin06

compressai_vision/model_wrappers/yolox.py

@@ -161,7 +161,13 @@ def input_to_features(self, x, device: str) -> Dict:
         """Computes deep features at the intermediate layer(s) all the way from the input"""
 
         self.model = self.model.to(device).eval()
-        img = x[0]["image"].unsqueeze(0).to(device)
+
+        # temp. solution
+        if not isinstance(x, torch.Tensor):
+            img = x[0]["image"].unsqueeze(0).to(device)
+        else:
+            img = x
+
         input_size = tuple(img.shape[2:])
 
         if self.split_id == self.SPLIT_L13: