X-iZhang commited on
Commit
5ccf46e
·
verified ·
1 Parent(s): bcc2db4

Update libra/eval/run_libra.py

Browse files
Files changed (1) hide show
  1. libra/eval/run_libra.py +3 -3
libra/eval/run_libra.py CHANGED
@@ -77,7 +77,7 @@ def load_images(image_file):
77
 
78
  return image
79
 
80
- def get_image_tensors(image_path, image_processor, model, device='cuda'):
81
  # Load and preprocess the images
82
  if isinstance(image_path, str):
83
  image = []
@@ -151,7 +151,7 @@ def libra_eval(
151
  conv.append_message(conv.roles[1], None)
152
  prompt = conv.get_prompt()
153
 
154
- input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).cuda()
155
  attention_mask = torch.ones(input_ids.shape, dtype=torch.long)
156
  pad_token_id = tokenizer.pad_token_id
157
 
@@ -162,7 +162,7 @@ def libra_eval(
162
  stopping_criteria = KeywordsStoppingCriteria(keywords, tokenizer, input_ids)
163
 
164
  with torch.inference_mode():
165
- torch.cuda.empty_cache()
166
  if num_beams > 1:
167
  output_ids = model.generate(
168
  input_ids=input_ids,
 
77
 
78
  return image
79
 
80
+ def get_image_tensors(image_path, image_processor, model, device='cpu'):
81
  # Load and preprocess the images
82
  if isinstance(image_path, str):
83
  image = []
 
151
  conv.append_message(conv.roles[1], None)
152
  prompt = conv.get_prompt()
153
 
154
+ input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).to("cpu")
155
  attention_mask = torch.ones(input_ids.shape, dtype=torch.long)
156
  pad_token_id = tokenizer.pad_token_id
157
 
 
162
  stopping_criteria = KeywordsStoppingCriteria(keywords, tokenizer, input_ids)
163
 
164
  with torch.inference_mode():
165
+
166
  if num_beams > 1:
167
  output_ids = model.generate(
168
  input_ids=input_ids,