diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py
index 03f2c44701..e1c43b8b85 100644
--- a/unsloth/models/rl.py
+++ b/unsloth/models/rl.py
@@ -900,9 +900,9 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"):
     if "temperature" in call_args:
         check_temperature = (
             "if temperature <= 0:\n"
-            "    raise MathError('Unsloth: Please set a positive non-zero temperature since your results will be wrong.')\n"
+            "    raise ValueError('Unsloth: Please set a positive non-zero temperature since your results will be wrong.')\n"
             "elif temperature >= 10:\n"
-            "    raise MathError('Unsloth: Please set a positive non-zero temperature less than 10, since sampling will be quite erratic.')\n"
+            "    raise ValueError('Unsloth: Please set a positive non-zero temperature less than 10, since sampling will be quite erratic.')\n"
             "\n"
         )
         extra_args += check_temperature
diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py
index 3dfeea6871..5e079335ae 100644
--- a/unsloth/models/rl_replacements.py
+++ b/unsloth/models/rl_replacements.py
@@ -861,13 +861,19 @@ def grpo_trainer_compute_loss(function_name, function):
                 else torch.tensor(0.0, device = self.model.device)
             )
             self._metrics[mode]["sampling/importance_sampling_ratio/min"].append(
-                nanmin(self.accelerator.gather(min_importance_sampling_ratio)).item()
+                self.accelerator.gather(min_importance_sampling_ratio)
+                .nan_to_num(nan = float("inf"))
+                .min()
+                .item()
             )
             self._metrics[mode]["sampling/importance_sampling_ratio/mean"].append(
                 self.accelerator.gather(mean_importance_sampling_ratio).nanmean().item()
             )
             self._metrics[mode]["sampling/importance_sampling_ratio/max"].append(
-                nanmax(self.accelerator.gather(max_importance_sampling_ratio)).item()
+                self.accelerator.gather(max_importance_sampling_ratio)
+                .nan_to_num(nan = float("-inf"))
+                .max()
+                .item()
             )
 
         return loss
diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py
index 36cfbf0b17..c909f963b9 100644
--- a/unsloth/models/vision.py
+++ b/unsloth/models/vision.py
@@ -147,7 +147,7 @@ def unsloth_base_fast_generate(
     elif "input_ids" in kwargs:
         input_ids = kwargs["input_ids"]
     elif "input" in kwargs:
-        input_ids = kwargs["input_ids"]
+        input_ids = kwargs["input"]
     elif "input_features" in kwargs:
         input_ids = kwargs["input_features"]
     elif "input_embeds" in kwargs:
@@ -156,7 +156,7 @@ def unsloth_base_fast_generate(
         input_ids = kwargs["inputs"]
     else:
         key = next(iter(kwargs.keys()))
-        if type(kwargs["key"]) is not torch.Tensor:
+        if type(kwargs[key]) is not torch.Tensor:
             raise TypeError("Unsloth: You need to pass in input_ids to .generate!")
         input_ids = kwargs[key]
     assert type(input_ids) is torch.Tensor