diff --git a/run.py b/run.py
index 0efb0e4d4de0799ac9cc02bf36f25d986a1a611d..78c18d1b8e2f5657d939d0dd825bc46fdb84af38 100644
--- a/run.py
+++ b/run.py
@@ -4,7 +4,8 @@ __date__ = '2019-11-14'
 
 import argparse
 
-from src.run_modules.experiment_setup import ExperimentSetup, PartitionCheck
+from src.run_modules.experiment_setup import ExperimentSetup
+from src.run_modules.partition_check import PartitionCheck
 from src.run_modules.model_setup import ModelSetup
 from src.run_modules.post_processing import PostProcessing
 from src.run_modules.pre_processing import PreProcessing
diff --git a/src/run_modules/experiment_setup.py b/src/run_modules/experiment_setup.py
index 45778074789683c612dfcf4a2d6a8b58315b2ad5..ac6806ebb80aafc179048b21bae323e035cc4b62 100644
--- a/src/run_modules/experiment_setup.py
+++ b/src/run_modules/experiment_setup.py
@@ -187,21 +187,6 @@ class ExperimentSetup(RunEnvironment):
             self._set_param("statistics_per_var", stat_new)
 
 
-class PartitionCheck(RunEnvironment):
-
-    """ Checking if running on a HPC login node. The onöy reason to run on login nodes is to download data. Training and validation should happen on compute nodes"""
-
-    def __init__(self):
-        # create run framework
-        super().__init__()
-
-        self._run()
-
-    def _run(self):
-        if self.data_store.get('hostname')[:2] in self.data_store.get('login_nodes'):
-            raise OSError('You are on a login node to download data. Use compute nodes and run again if you want to train and validate a model.')
-
-
 if __name__ == "__main__":
 
     formatter = '%(asctime)s - %(levelname)s: %(message)s  [%(filename)s:%(funcName)s:%(lineno)s]'
diff --git a/src/run_modules/partition_check.py b/src/run_modules/partition_check.py
index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..ee4edf0523000cd61d308e1f5e10b1f604029b19 100644
--- a/src/run_modules/partition_check.py
+++ b/src/run_modules/partition_check.py
@@ -0,0 +1,23 @@
+__author__ = "Felix Kleinert"
+__date__ = '2020-04-07'
+
+from src.run_modules.run_environment import RunEnvironment
+
+
+class PartitionCheck(RunEnvironment):
+    """
+    Checking if running on a HPC login node. The only reason to run on login nodes is to download data.
+    Training and validation should happen on compute nodes
+    """
+
+    def __init__(self):
+        # create run framework
+        super().__init__()
+
+        self._run()
+
+    def _run(self):
+        if self.data_store.get('hostname')[:2] in self.data_store.get('login_nodes'):
+            raise OSError(
+                'You are on a login node to download data. Use compute nodes and run again if '
+                'you want to train and validate a model.')