diff --git a/above_shrubs/pipelines/chm_pipeline.py b/above_shrubs/pipelines/chm_pipeline.py
index 6ed72a9..97ba102 100644
--- a/above_shrubs/pipelines/chm_pipeline.py
+++ b/above_shrubs/pipelines/chm_pipeline.py
@@ -144,14 +144,14 @@ def train(self):
         max_epochs = 50
         fast_dev_run = False
 
-        datamodule = CHMDataModule(
-            train_data_dir=self.conf.train_data_dir,
-            train_label_dir=self.conf.train_label_dir,
-            test_data_dir=self.conf.test_data_dir,
-            test_label_dir=self.conf.test_label_dir,
-            batch_size=16,
-            num_workers=8,
-        )
+        #datamodule = CHMDataModule(
+        #    train_data_dir=self.conf.train_data_dir,
+        #    train_label_dir=self.conf.train_label_dir,
+        #    test_data_dir=self.conf.test_data_dir,
+        #    test_label_dir=self.conf.test_label_dir,
+        #    batch_size=16,
+        #    num_workers=8,
+        #)
 
         # Set main data loader
         #chm_train_dataset = CHMDataset(
diff --git a/requirements/Dockerfile b/requirements/Dockerfile
index f3c5426..adc86c5 100644
--- a/requirements/Dockerfile
+++ b/requirements/Dockerfile
@@ -1,5 +1,5 @@
 # Arguments to pass to the image
-ARG VERSION_DATE=24.09
+ARG VERSION_DATE=24.01
 ARG FROM_IMAGE=nvcr.io/nvidia/pytorch
 
 # Import RAPIDS container as the BASE Image (cuda base image)
@@ -17,7 +17,7 @@ RUN apt-get update && \
 # Pip
 RUN pip --no-cache-dir install --ignore-installed omegaconf \
     #terratorch \
-    pytorch-lightning \
+    #pytorch-lightning \
     #Lightning \
     transformers \
     datasets \