bug fix

FoundationVision · Jun 7, 2024 · d3a95b6 · d3a95b6
1 parent a211671
commit d3a95b6
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/groma/eval/run_groma.py b/groma/eval/run_groma.py
@@ -55,7 +55,10 @@ def eval_model(model_name, quant_type, image_file, query,):
         )
         kwargs = {'quantization_config': int4_quant_cfg}
 
-    model = GromaModel.from_pretrained(model_name, **kwargs).cuda()        
+    if quant_type == '8bit' or quant_type == '4bit':
+        model = GromaModel.from_pretrained(model_name, **kwargs)
+    else:
+        model = GromaModel.from_pretrained(model_name, **kwargs).cuda()
     model.init_special_token_id(tokenizer)
 
     conversations = []