diff --git a/.gitignore b/.gitignore
index ebbac55..90cdd97 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,3 +1,6 @@
+*.err
+*.out
+*.log
 output
 # Created by https://www.toptal.com/developers/gitignore/api/python,node,git,visualstudiocode
 # Edit at https://www.toptal.com/developers/gitignore?templates=python,node,git,visualstudiocode
diff --git a/aimodel/slurm-pretrain.job b/aimodel/slurm-pretrain.job
index b559e0a..6380206 100755
--- a/aimodel/slurm-pretrain.job
+++ b/aimodel/slurm-pretrain.job
@@ -3,8 +3,8 @@
 #SBATCH -N 1
 #SBATCH -n 4
 #SBATCH --gres=gpu:1
-#SBATCH -o %j.%N.%a.out
-#SBATCH -e %j.%N.%a.err
+#SBATCH -o %j.%N.%a.out.log
+#SBATCH -e %j.%N.%a.err.log
 #SBATCH -p gpu05
 #SBATCH --time=5-00:00:00
 #SBATCH --exclusive
diff --git a/aimodel/src/lib/ai/model_rainfallwater_contrastive.py b/aimodel/src/lib/ai/model_rainfallwater_contrastive.py
index f7af3a7..f32fd51 100644
--- a/aimodel/src/lib/ai/model_rainfallwater_contrastive.py
+++ b/aimodel/src/lib/ai/model_rainfallwater_contrastive.py
@@ -9,8 +9,11 @@ from .components.LossContrastive import LossContrastive
 def model_rainfallwater_contrastive(shape_rainfall, shape_water, feature_dim=200):
 	logger.info(shape_rainfall)
 	logger.info(shape_water)
-	rainfall_width, rainfall_height, rainfall_channels = shape_rainfall
-	water_width, water_height, water_channels = shape_water
+	
+	# Shapes come from what rainfallwrangler sees them as, but we add an extra dimension when reading the .tfrecord file
+	rainfall_width, rainfall_height, rainfall_channels = shape_rainfall # shape = [width, height, channels]
+	water_width, water_height = shape_water # shape = [width, height]
+	water_channels = 1 # added in dataset → make_dataset → parse_item
 	
 	input_rainfall = tf.keras.layers.Input(
 		shape=shape_rainfall
diff --git a/aimodel/src/lib/dataset/dataset.py b/aimodel/src/lib/dataset/dataset.py
index 3fac94d..1435e7c 100644
--- a/aimodel/src/lib/dataset/dataset.py
+++ b/aimodel/src/lib/dataset/dataset.py
@@ -20,6 +20,9 @@ def parse_item(item):
 	rainfall = tf.io.parse_tensor(parsed["rainfallradar"], out_type=tf.float32)
 	water = tf.io.parse_tensor(parsed["waterdepth"], out_type=tf.float32)
 	
+	# [width, height] → [width, height, channels]
+	water = tf.expand_dims(water, axis=-1)
+	
 	# TODO: The shape of the resulting tensor can't be statically determined, so we need to reshape here
 	
 	# TODO: Any other additional parsing here, since multiple .map() calls are not optimal