diff --git a/astrai/parallel/setup.py b/astrai/parallel/setup.py
index 2ce838c..e8d931b 100644
--- a/astrai/parallel/setup.py
+++ b/astrai/parallel/setup.py
@@ -12,7 +12,7 @@ import torch
 import torch.distributed as dist
 import torch.multiprocessing as mp
 
-from astrai.parallel.signal_handler import install_early_signal_handlers
+from astrai.signal_handler import install_early_signal_handlers
 
 logger = logging.getLogger(__name__)
 
diff --git a/astrai/parallel/signal_handler.py b/astrai/signal_handler.py
similarity index 100%
rename from astrai/parallel/signal_handler.py
rename to astrai/signal_handler.py
diff --git a/astrai/trainer/trainer.py b/astrai/trainer/trainer.py
index 6292081..fa09bc8 100644
--- a/astrai/trainer/trainer.py
+++ b/astrai/trainer/trainer.py
@@ -5,7 +5,7 @@ import torch.distributed as dist
 
 from astrai.config import TrainConfig
 from astrai.parallel.setup import spawn_parallel_fn
-from astrai.parallel.signal_handler import (
+from astrai.signal_handler import (
     register_signal_handlers,
     unregister_signal_handlers,
 )
diff --git a/tests/trainer/test_signal_handler.py b/tests/trainer/test_signal_handler.py
index d1750d3..ece946b 100644
--- a/tests/trainer/test_signal_handler.py
+++ b/tests/trainer/test_signal_handler.py
@@ -10,7 +10,7 @@ from torch.utils.data import Dataset
 
 from astrai.config import TrainConfig
 from astrai.model.transformer import AutoRegressiveLM
-from astrai.parallel.signal_handler import register_signal_handlers
+from astrai.signal_handler import register_signal_handlers
 from astrai.trainer import Trainer
 from astrai.trainer.schedule import SchedulerFactory
 from astrai.trainer.train_context import TrainContext