Update inference/infer.py

inference/infer.py

@@ -69,7 +69,12 @@ 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 = "cuda"
+
+# Check if CUDA is available
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+# Now you can use `device` to move your tensors or models to the GPU (if available)
+print(f"Using device: {device}")
+
 mmtokenizer = _MMSentencePieceTokenizer("./mm_tokenizer_v0.2_hf/tokenizer.model")
 model = AutoModelForCausalLM.from_pretrained(
     stage1_model,