fix: add torch_dtype=bfloat16 to prevent Float/BFloat16 mismatch in GRPO

train/grpo_train.py

@@ -115,6 +115,7 @@ def _load_model_and_tokenizer(model_id: str, seed: int):
     model = AutoModelForCausalLM.from_pretrained(
         model_id,
         quantization_config=bnb,
+        torch_dtype=torch.bfloat16,
         device_map="auto",
         trust_remote_code=True,
     )