kastan commited on
Commit
421e169
β€’
1 Parent(s): f38e35a

remove fancy CUDA device selection, for plain 'cuda'

Browse files
Files changed (1) hide show
  1. retrieval.py +5 -5
retrieval.py CHANGED
@@ -9,12 +9,12 @@ from langchain.vectorstores import Pinecone
9
 
10
  # load custom code
11
  from clip_for_ppts import ClipImage
12
- from gpu_memory_utils import (get_device_with_most_free_memory,
13
- get_free_memory_dict,
14
- get_gpu_ids_with_sufficient_memory)
15
 
16
- # from PIL import Image
 
 
17
 
 
18
 
19
  LECTURE_SLIDES_DIR = os.path.join(os.getcwd(), "lecture_slides")
20
  PINECONE_API_KEY = os.environ.get("PINECONE_API")
@@ -81,7 +81,7 @@ class Retrieval:
81
  self.clip_search_class = ClipImage(path_of_ppt_folders=LECTURE_SLIDES_DIR,
82
  path_to_save_image_features=os.getcwd(),
83
  mode='text',
84
- device=f'cuda:{get_device_with_most_free_memory()}')
85
 
86
  def reverse_img_search(self, img):
87
  imgs = self.clip_search_class.image_to_images_search(img)
 
9
 
10
  # load custom code
11
  from clip_for_ppts import ClipImage
 
 
 
12
 
13
+ # from gpu_memory_utils import (get_device_with_most_free_memory,
14
+ # get_free_memory_dict,
15
+ # get_gpu_ids_with_sufficient_memory)
16
 
17
+ # from PIL import Image
18
 
19
  LECTURE_SLIDES_DIR = os.path.join(os.getcwd(), "lecture_slides")
20
  PINECONE_API_KEY = os.environ.get("PINECONE_API")
 
81
  self.clip_search_class = ClipImage(path_of_ppt_folders=LECTURE_SLIDES_DIR,
82
  path_to_save_image_features=os.getcwd(),
83
  mode='text',
84
+ device='cuda')
85
 
86
  def reverse_img_search(self, img):
87
  imgs = self.clip_search_class.image_to_images_search(img)