Fix multiple GPU inference.

dr-data · Nov 8, 2023 · 4a77fb4 · 4a77fb4
1 parent a7f6cce
commit 4a77fb4
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/llava/model/llava_arch.py b/llava/model/llava_arch.py
@@ -116,9 +116,9 @@ def prepare_inputs_labels_for_multimodal(
             image_features = self.encode_images(concat_images)
             split_sizes = [image.shape[0] for image in images]
             image_features = torch.split(image_features, split_sizes, dim=0)
-            image_features = [x.flatten(0, 1).to(concat_images.device) for x in image_features]
+            image_features = [x.flatten(0, 1).to(self.device) for x in image_features]
         else:
-            image_features = self.encode_images(images)
+            image_features = self.encode_images(images).to(self.device)
 
         # TODO: image start / end is not implemented here to support pretraining.
         if getattr(self.config, 'tune_mm_mlp_adapter', False) and getattr(self.config, 'mm_use_im_start_end', False):