From e88f8283a9f85cb3575dd8a3d9257f23323e65ff Mon Sep 17 00:00:00 2001
From: root <403644786@qq.com>
Date: Tue, 30 Jul 2024 14:44:56 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E6=94=B9=E4=BA=86=E4=B8=80=E4=B8=8B?=
 =?UTF-8?q?=E4=BD=BF=E7=94=A8cuda=E7=9A=84=E6=96=AD=E8=A8=80?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 quantize/bnb_quantize.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/quantize/bnb_quantize.py b/quantize/bnb_quantize.py
index 205cd28..7aa7b46 100644
--- a/quantize/bnb_quantize.py
+++ b/quantize/bnb_quantize.py
@@ -18,11 +18,14 @@ import torch
 import GPUtil
 import os
 
+assert torch.cuda.is_available(),"CUDA is not available, but this code requires a GPU."
+
+device = 'cuda'  # Select GPU to use
 model_path = '/root/ld/ld_model_pretrained/MiniCPM-Llama3-V-2_5' # Model download path
-device = 'cuda' # Select GPU if available, otherwise CPU
 save_path = '/root/ld/ld_model_pretrain/MiniCPM-Llama3-V-2_5_int4' # Quantized model save path
 image_path = './assets/airplane.jpeg'
 
+
 # Create a configuration object to specify quantization parameters
 quantization_config = BitsAndBytesConfig(
     load_in_4bit=True,  # Whether to perform 4-bit quantization