update monkey.py to save GPU mem

mrG7 · Jan 10, 2024 · 2f11105 · 2f11105
1 parent 37aec91
commit 2f11105
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/vlmeval/vlm/monkey.py b/vlmeval/vlm/monkey.py
@@ -13,7 +13,9 @@ def __init__(self, model_path='echo840/Monkey', **kwargs):
         assert model_path is not None
         self.model_path = model_path
         self.tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
-        self.model = AutoModelForCausalLM.from_pretrained(model_path, device_map='cuda', trust_remote_code=True).eval()
+        model = AutoModelForCausalLM.from_pretrained(model_path, device_map='cpu', trust_remote_code=True)
+        model.eval()
+        self.model = model.cuda()
         self.kwargs = kwargs
         warnings.warn(f"Following kwargs received: {self.kwargs}, will use as generation config. ")
         torch.cuda.empty_cache()