diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py
index 855bed66c0..b5fd37056c 100644
--- a/unsloth/models/llama.py
+++ b/unsloth/models/llama.py
@@ -1385,14 +1385,16 @@ class FastLlamaModel:
         # Now patch lm_head and embed_tokens
         if train_embed_tokens:
             print("Unsloth: Casting embed_tokens to float32")
-            model.model.model.embed_tokens.to(torch.float32)
-            model.model.model.embed_tokens.requires_grad_(True)
+            assert(hasattr(model.model.model.embed_tokens, "modules_to_save"))
+            model.model.model.embed_tokens.modules_to_save.default.to(torch.float32)
+            model.model.model.embed_tokens.modules_to_save.default.requires_grad_(True)
         pass
 
         if train_lm_head:
             print("Unsloth: Casting lm_head to float32")
-            model.model.lm_head.to(torch.float32)
-            model.model.lm_head.requires_grad_(True)
+            assert(hasattr(model.model.lm_head, "modules_to_save"))
+            model.model.lm_head.modules_to_save.default.to(torch.float32)
+            model.model.lm_head.modules_to_save.default.requires_grad_(True)
         pass
 
         return model