diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py
index 22fb51144c..b7333f0098 100644
--- a/unsloth/models/_utils.py
+++ b/unsloth/models/_utils.py
@@ -79,6 +79,7 @@ __all__ = [
     "offload_input_embeddings",
     "offload_output_embeddings",
     "is_bfloat16_supported",
+    "unsloth_offloaded_gradient_checkpoint",
 ]
 
 
@@ -402,6 +403,12 @@ class Unsloth_Offloaded_Gradient_Checkpointer(torch.autograd.Function):
 pass
 
 
+@torch._disable_dynamo
+def unsloth_offloaded_gradient_checkpoint(function, *args, use_reentrant = None, **kwargs):
+    return Unsloth_Offloaded_Gradient_Checkpointer.apply(function, *args)
+pass
+
+
 """
     Remove warnings about missing kwargs
 """
diff --git a/unsloth/tokenizer_utils.py b/unsloth/tokenizer_utils.py
index 03f3e341d7..6e4d69107e 100644
--- a/unsloth/tokenizer_utils.py
+++ b/unsloth/tokenizer_utils.py
@@ -577,7 +577,9 @@ def fix_untrained_tokens(model, tokenizer, train_dataset, eps = 1e-16):
     
     where_untrained_set = frozenset(where_untrained)
     actual_bad_tokens = tokenizer.convert_ids_to_tokens(where_untrained)
-
+    # Remove None items in actual_bad_tokens
+    actual_bad_tokens = [x for x in actual_bad_tokens if x is not None]
+    
     # Check if tokenizer and training datasets have bad tokens
     if_bad_first  = False
     if_bad_second = False