diff --git a/llms/mlx_lm/models/plamo2.py b/llms/mlx_lm/models/plamo2.py
index 0fbb4cee..c6e413b0 100644
--- a/llms/mlx_lm/models/plamo2.py
+++ b/llms/mlx_lm/models/plamo2.py
@@ -7,7 +7,7 @@ from typing import Any, Literal, NamedTuple, Optional, Union
 import mlx.core as mx
 import mlx.nn as nn
 
-from .base import BaseModelArgs, create_attention_mask
+from mlx_lm.models.base import BaseModelArgs, create_attention_mask
 
 
 def _is_first_token(mask: mx.array) -> mx.array: