Update inference/infer.py

inference/infer.py

@@ -68,7 +68,8 @@ os.makedirs(stage1_output_dir, exist_ok=True)
 os.makedirs(stage2_output_dir, exist_ok=True)
 
 # load tokenizer and model
-device = torch.device(f"cuda:{cuda_idx}" if torch.cuda.is_available() else "cpu")
+#device = torch.device(f"cuda:{cuda_idx}" if torch.cuda.is_available() else "cpu")
+device = "cuda"
 mmtokenizer = _MMSentencePieceTokenizer("./mm_tokenizer_v0.2_hf/tokenizer.model")
 model = AutoModelForCausalLM.from_pretrained(
     stage1_model,