diff --git a/src/rail/core/stage.py b/src/rail/core/stage.py
index 2576ee9..1394fef 100644
--- a/src/rail/core/stage.py
+++ b/src/rail/core/stage.py
@@ -5,7 +5,8 @@
 import yaml
 from math import ceil
 
-from ceci import PipelineStage, MiniPipeline
+from ceci.stage import PipelineStage
+from ceci.pipeline import MiniPipeline
 from ceci.config import StageParameter as Param
 from rail.core.data import DATA_STORE, DataHandle