Update README.md

README.md

@@ -26,10 +26,11 @@ The DiffuCoder-7B-Instruct model builds on the DiffuCoder-7B-Base checkpoint wit
 import torch
 from transformers import AutoModel, AutoTokenizer
 
+device = "cuda"
 model_path = "apple/DiffuCoder-7B-Instruct"
 model = AutoModel.from_pretrained(model_path, torch_dtype=torch.bfloat16, trust_remote_code=True)
 tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
-model = model.to("cuda").eval()
+model = model.to(device).eval()
 
 query = "Write a function to find the shared elements from the given two lists."
 prompt = f"""<|im_start|>system
@@ -43,8 +44,8 @@ You are a helpful assistant.<|im_end|>
 TOKEN_PER_STEP = 1 # diffusion timesteps * TOKEN_PER_STEP = total new tokens
 
 inputs = tokenizer(prompt, return_tensors="pt")
-input_ids = inputs.input_ids.to(device="cuda")
-attention_mask = inputs.attention_mask.to(device="cuda")
+input_ids = inputs.input_ids.to(device=device)
+attention_mask = inputs.attention_mask.to(device=device)
 
 output = model.diffusion_generate(
     input_ids,