Update inference/infer.py

inference/infer.py

@@ -78,7 +78,9 @@ model = AutoModelForCausalLM.from_pretrained(
     stage1_model, 
     torch_dtype=torch.bfloat16,
     attn_implementation="flash_attention_2", # To enable flashattn, you have to install flash-attn
-    ).to(device).eval()
+    )
+model.to(device)
+model.eval()
 
 codectool = CodecManipulator("xcodec", 0, 1)
 codectool_stage2 = CodecManipulator("xcodec", 0, 8)