diff --git a/unsloth/models/mistral.py b/unsloth/models/mistral.py
index 0eed45c5cd..5e893d2b6f 100644
--- a/unsloth/models/mistral.py
+++ b/unsloth/models/mistral.py
@@ -307,9 +307,9 @@ def MistralForCausalLM_fast_forward(
             RETURN_LOGITS = False
 
         if not RETURN_LOGITS and labels is not None:
-            n_items = kwargs.get("num_items_in_batch", None) or kwargs.get(
-                "n_items", None
-            )
+            n_items = kwargs.get("num_items_in_batch", None)
+            if n_items is None:
+                n_items = kwargs.get("n_items", None)
             logit_softcapping = getattr(self.config, "final_logit_softcapping", 0)
 
             # loss = fused_linear_cross_entropy(
@@ -363,11 +363,13 @@ def MistralForCausalLM_fast_forward(
             shift_labels,
             kwargs.get("packed_seq_lengths"),
         )
+        n_items = kwargs.get("num_items_in_batch", None)
+        if n_items is None:
+            n_items = kwargs.get("n_items", None)
         loss = fast_cross_entropy_loss(
             logits = shift_logits,
             labels = shift_labels,
-            n_items = kwargs.get("num_items_in_batch", None)
-            or kwargs.get("n_items", None),
+            n_items = n_items,
         )
 
     if not return_dict:
diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py
index 003a0e7f1b..03f2c44701 100644
--- a/unsloth/models/rl.py
+++ b/unsloth/models/rl.py
@@ -22,7 +22,6 @@ from typing import Any, Callable, Dict, List, Literal, Optional, Tuple, Union
 import inspect
 import os
 import re
-import torch
 from unsloth_zoo.compiler import create_new_function
 from unsloth_zoo.log import logger
 from unsloth_zoo.logging_utils import PatchRLStatistics
diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py
index e1cf8f6f82..36cfbf0b17 100644
--- a/unsloth/models/vision.py
+++ b/unsloth/models/vision.py
@@ -68,11 +68,9 @@ import functools
 import os
 import gc
 import math
-import functools
 from typing import Optional, Tuple, List, Union
 import re, inspect, sys
 import contextlib
-import types
 
 try:
     from huggingface_hub.utils import get_token