diff --git a/tests/model/test_aiia.py b/tests/model/test_aiia.py
index 340537b..bffa616 100644
--- a/tests/model/test_aiia.py
+++ b/tests/model/test_aiia.py
@@ -1,6 +1,6 @@
 import os
 import torch
-from aiia import AIIABase, AIIABaseShared, AIIAExpert, AIIAmoe, AIIAchunked, AIIAConfig, AIIASparseMoe
+from aiia.model.Model import AIIABase, AIIABaseShared, AIIAExpert, AIIAmoe, AIIAchunked, AIIAConfig, AIIASparseMoe
 
 def test_aiiabase_creation():
     config = AIIAConfig()