From a9d890cc448fdd7420c59b002cd8beaa6cbe8464 Mon Sep 17 00:00:00 2001 From: chflame163 Date: Thu, 6 Mar 2025 17:26:54 +0800 Subject: [PATCH] fix load model to cuda bug --- py/cogview4.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/py/cogview4.py b/py/cogview4.py index f3e8c1f..23900a1 100644 --- a/py/cogview4.py +++ b/py/cogview4.py @@ -61,11 +61,11 @@ class CogView4_Wrapper: if self.dtype != dtype or self.model_name != model: model_path = check_and_download_model(model, f"THUDM/{model}") if dtype == 'bf16': - self.pipe = CogView4Pipeline.from_pretrained(model_path, torch_dtype=torch.bfloat16).to("cuda") + self.pipe = CogView4Pipeline.from_pretrained(model_path, torch_dtype=torch.bfloat16) self.model_name = model self.dtype = dtype else: - self.pipe = CogView4Pipeline.from_pretrained(model_path, torch_dtype=torch.float32).to("cuda") + self.pipe = CogView4Pipeline.from_pretrained(model_path, torch_dtype=torch.float32) self.model_name = model self.dtype = dtype # Open it for reduce GPU memory usage