diff --git a/src/mflux/run.py b/src/mflux/run.py
index facae1c..4be1580 100644
--- a/src/mflux/run.py
+++ b/src/mflux/run.py
@@ -7,7 +7,7 @@ sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), 'src')))
 
 from mflux.config.model_config import ModelConfig
 from mflux.config.config import Config
-from mflux.flux import Flux1
+from mflux.flux.flux import Flux1
 
 
 def main():
diff --git a/src/mflux/save.py b/src/mflux/save.py
index d38f19c..304a1e1 100644
--- a/src/mflux/save.py
+++ b/src/mflux/save.py
@@ -4,7 +4,7 @@ import sys
 
 sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), 'src')))
 
-from mflux.flux import Flux1
+from mflux.flux.flux import Flux1
 from mflux.config.model_config import ModelConfig