enhanced Visual Transformers batch extraction.

sisinflab · Mar 28, 2024 · e3b031c · e3b031c
1 parent f79015d
commit e3b031c
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/ducho/multimodal/visual/VisualFeatureExtractor.py b/ducho/multimodal/visual/VisualFeatureExtractor.py
@@ -104,7 +104,7 @@ def extract_feature(self, image):
 
         elif 'transformers' in self._backend_libraries_list:
             # converting the input image tensor - outcome of the pre-processor - in a set.
-            model_input = {'pixel_values': image}
+            model_input = {'pixel_values': image[0]}
             model_input = {k: torch.tensor(v).to(self._device) for k, v in model_input.items()}
             model_output = getattr(self._model(**model_input), self._output_layer.lower())
             return model_output.detach().cpu().numpy()