diff --git a/loader.py b/loader.py
index af86bd3..992aab9 100644
--- a/loader.py
+++ b/loader.py
@@ -242,10 +242,8 @@ def gguf_clip_loader(path):
             sd[temb_key] = dequantize_tensor(sd[temb_key], dtype=torch.float16)
         sd = sd_map_replace(sd, T5_SD_MAP)
     elif arch in {"llama"}:
+        # TODO: pass model_options["vocab_size"] to loader somehow
         temb_key = "token_embd.weight"
-        if temb_key in sd and sd[temb_key].shape != (128320, 4096):
-            # This still works. TODO: Only relevant for llava?
-            logging.warning("Warning! token_embd shape may be incorrect for llama 3 model!")
         if temb_key in sd and sd[temb_key].shape[0] >= (64 * 1024):
             # See note above for T5.
             logging.warning(f"Dequantizing {temb_key} to prevent runtime OOM.")
diff --git a/ops.py b/ops.py
index 0a04164..441b301 100644
--- a/ops.py
+++ b/ops.py
@@ -111,6 +111,9 @@ class GGMLLayer(torch.nn.Module):
         # NOTE: using modified load for linear due to not initializing on creation, see GGMLOps todo
         if self.is_ggml_quantized(weight=weight, bias=bias) or isinstance(self, torch.nn.Linear):
             return self.ggml_load_from_state_dict(state_dict, prefix, *args, **kwargs)
+        # Not strictly required, but fixes embedding shape mismatch. Threshold set in loader.py
+        if isinstance(self, torch.nn.Embedding) and self.weight.shape[0] >= (64 * 1024):
+            return self.ggml_load_from_state_dict(state_dict, prefix, *args, **kwargs)
         return super()._load_from_state_dict(state_dict, prefix, *args, **kwargs)
 
     def ggml_load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs):