Merge pull request #3986 from voxel51/open_clip_gpu_support

Adding gpu support for openclip models
voxel51 · Jan 12, 2024 · cf4659c · cf4659c
2 parents 861913b + c7f1590
commit cf4659c
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/fiftyone/utils/open_clip.py b/fiftyone/utils/open_clip.py
@@ -90,7 +90,9 @@ def _load_model(self, config):
             _,
             self.preprocess,
         ) = open_clip.create_model_and_transforms(
-            config.clip_model, pretrained=config.pretrained
+            config.clip_model,
+            pretrained=config.pretrained,
+            device=config.device,
         )
         self._tokenizer = open_clip.get_tokenizer(config.clip_model)
         return self._model
@@ -102,6 +104,8 @@ def _get_text_features(self):
             ]
             # Tokenize text
             text = self._tokenizer(prompts)
+            if self._using_gpu:
+                text = text.cuda()
             self._text_features = self._model.encode_text(text)
 
         return self._text_features