Update README.md

README.md

@@ -30,6 +30,9 @@ Before running the snippet, you need to install the following dependencies:
 pip install torch transformers accelerate pillow
 ```
 
+If the CUDA memory is enough, it would be faster to execute this snippet by setting `CUDA_VISIBLE_DEVICES=0`.
+
+
 ```python
 import torch
 import transformers
@@ -43,12 +46,13 @@ transformers.logging.disable_progress_bar()
 warnings.filterwarnings('ignore')
 
 # set device
-torch.set_default_device('cpu')  # or 'cuda'
+device = 'cuda'  # or cpu
+torch.set_default_device(device)
 
 # create model
 model = AutoModelForCausalLM.from_pretrained(
     'BAAI/Bunny-v1_0-3B',
-    torch_dtype=torch.float16,
+    torch_dtype=torch.float16, # float32 for cpu
     device_map='auto',
     trust_remote_code=True)
 tokenizer = AutoTokenizer.from_pretrained(
@@ -59,11 +63,11 @@ tokenizer = AutoTokenizer.from_pretrained(
 prompt = 'Why is the image funny?'
 text = f"A chat between a curious user and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the user's questions. USER: <image>\n{prompt} ASSISTANT:"
 text_chunks = [tokenizer(chunk).input_ids for chunk in text.split('<image>')]
-input_ids = torch.tensor(text_chunks[0] + [-200] + text_chunks[1], dtype=torch.long).unsqueeze(0)
+input_ids = torch.tensor(text_chunks[0] + [-200] + text_chunks[1], dtype=torch.long).unsqueeze(0).to(device)
 
 # image, sample images can be found in images folder
 image = Image.open('example_2.png')
-image_tensor = model.process_images([image], model.config).to(dtype=model.dtype)
+image_tensor = model.process_images([image], model.config).to(dtype=model.dtype, device=device)
 
 # generate
 output_ids = model.generate(