diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml
index 5b522c26..cc2fc6c4 100644
--- a/.github/workflows/test.yaml
+++ b/.github/workflows/test.yaml
@@ -28,7 +28,9 @@ jobs:
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
-          pip install -r requirements/required.txt -r requirements/test.txt -r requirements/optional.txt
+          #pip install -r requirements/required.txt -r requirements/test.txt -r requirements/optional.txt
+          pip install -r requirements/test.txt -r requirements/optional.txt
+          pip install -e .[torchgeo]
           pip install git+https://github.com/NASA-IMPACT/Prithvi-WxC.git
           pip install git+https://github.com/IBM/granite-wxc.git
       - name: List pip dependencies
diff --git a/mkdocs.yml b/mkdocs.yml
index 3105bbf0..0ffd9b8b 100644
--- a/mkdocs.yml
+++ b/mkdocs.yml
@@ -15,9 +15,9 @@ plugins:
           paths: [src] # search packages in the src folde
           options:
             show_root_heading: true
-extra:
-  version:
-    provider: mike
+              #extra:
+              #  version:
+              #    provider: mike
 
 site_url: https://ibm.github.io/terratorch/
 repo_url: https://github.com/IBM/terratorch
diff --git a/pyproject.toml b/pyproject.toml
index 6a477ac4..30d3cc0a 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -34,7 +34,7 @@ dependencies = [
   "torchgeo>=0.6.0",
   "rioxarray>=0.15.0",
   # see issue #64
-  "albumentations>=1.3.1, <=1.4.10",
+  "albumentations>=1.3.1, <=1.4.21",
   "albucore<=0.0.16",
   "rasterio>=1.3.9",
   "torchmetrics<=1.3.1",
@@ -44,11 +44,29 @@ dependencies = [
   "mlflow>=2.12.1",
   # broken due to https://github.com/Lightning-AI/pytorch-lightning/issues/19977
   "lightning[pytorch-extra]>=2,!=2.3.*",
-  "segmentation-models-pytorch>=0.3"
+  "segmentation-models-pytorch>=0.3",
+  "jsonargparse<=4.35.0",
   # Dependencies not available on PyPI
 ]
 
 [project.optional-dependencies]
+torchgeo = [
+  "torch==2.4.1",
+  "torchvision==0.19.1",
+  "torchgeo @ git+https://github.com/microsoft/torchgeo.git@fedf99375535f801565856cd774bfa9e5a251d55",
+  "rioxarray>=0.15.0",
+  "albumentations==1.3.1",
+  "albucore<=0.0.16",
+  "rasterio>=1.3.9",
+  "torchmetrics<=1.3.1",
+  "geopandas>=0.14.4",
+  "lightly>=1.4.25",
+  "h5py>=3.10.0",
+  "mlflow>=2.12.1",
+  "lightning[pytorch-extra]>=2,!=2.3.*",
+  "segmentation-models-pytorch>=0.3"
+]
+
 dev = [
   "black",
   "mkdocs-material",