diff --git a/src/twomartens/masterthesis/aae/train.py b/src/twomartens/masterthesis/aae/train.py
index b288d10..2f7d700 100644
--- a/src/twomartens/masterthesis/aae/train.py
+++ b/src/twomartens/masterthesis/aae/train.py
@@ -127,8 +127,10 @@ def train_mnist(folding_id: int, inlier_classes: Sequence[int], total_classes: i
     total_lowest_loss = math.inf
     grace_period = GRACE
     
-    checkpoint_prefix = './weights/' + str(inlier_classes[0]) + '/' + str(iteration) + '/ckpt'
-    latest_checkpoint = tf.train.latest_checkpoint(checkpoint_prefix)
+    checkpoint_dir = os.path.abspath('./weights/' + str(inlier_classes[0]) + '/' + str(iteration) + '/')
+    os.makedirs(checkpoint_dir, exist_ok=True)
+    checkpoint_prefix = os.path.join(checkpoint_dir, 'ckpt')
+    latest_checkpoint = tf.train.latest_checkpoint(checkpoint_dir)
     checkpoint = tf.train.Checkpoint(encoder=encoder,
                                      decoder=decoder,
                                      z_discriminator=z_discriminator,