diff --git a/ANN_Training.py b/ANN_Training.py
index 34cca6cfdc5c8c0797a825f86d6b7ce7e9d733f2..b52964c379df236e836be2fe97f0522dc9acd2ef 100644
--- a/ANN_Training.py
+++ b/ANN_Training.py
@@ -2,9 +2,11 @@
 """
 @author: Laura C. Kühle, Soraya Terrab (sorayaterrab)
 
+TODO: Test new ANN set-up with Soraya
+TODO: Integrate Main.py into Snakefile
 TODO: Add log to pipeline -> Done
 TODO: Remove object set-up (for more flexibility)
-TODO: Adapt TCD file to new classification
+TODO: Adapt TCD file to new classification -> Done
 TODO: Add documentation
 TODO: Improve log output
 TODO: Remove unnecessary comments -> Done
diff --git a/Troubled_Cell_Detector.py b/Troubled_Cell_Detector.py
index 74b75e1c0d7b5bfbf4e2f965e2d9880359ba5024..ef91fcc84c81aea9ccc9ef2ac9173afa76d7da66 100644
--- a/Troubled_Cell_Detector.py
+++ b/Troubled_Cell_Detector.py
@@ -265,7 +265,7 @@ class ArtificialNeuralNetwork(TroubledCellDetector):
             'input_size': self._stencil_len+2, 'first_hidden_size': 8, 'second_hidden_size': 4,
             'output_size': 2, 'activation_function': 'Softmax', 'activation_config': {'dim': 1}})
         self._model_state = config.pop(
-            'model_state', 'Train24k24k_Valid8k8k_Norm12ReLU8+4nodesSM1Adamlr1e-2MSE.pt')
+            'model_state', 'Snakemake-Test/trained models/model__Adam.pt')
 
         if not hasattr(ANN_Model, self._model):
             raise ValueError('Invalid model: "%s"' % self._model)