From 787099d1eaa45afc2c3007a7919c431737310ec3 Mon Sep 17 00:00:00 2001
From: Felix Stollenwerk <felix.stollenwerk@ai.se>
Date: Wed, 6 Nov 2024 18:38:01 +0100
Subject: [PATCH 1/2] fix: import T_co from torch

---
 src/modalities/dataloader/dataloader.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/src/modalities/dataloader/dataloader.py b/src/modalities/dataloader/dataloader.py
index fbf5cc36..ea01f21e 100644
--- a/src/modalities/dataloader/dataloader.py
+++ b/src/modalities/dataloader/dataloader.py
@@ -1,7 +1,12 @@
 from typing import Iterable, Optional, Union
 
 from torch.utils.data import Dataset, DistributedSampler, Sampler
-from torch.utils.data.dataloader import DataLoader, T_co, _collate_fn_t, _worker_init_fn_t
+from torch.utils.data.dataloader import DataLoader, _collate_fn_t, _worker_init_fn_t
+
+try:  # torch <= 2.4
+    from torch.utils.data.dataloader import T_co
+except ImportError:  # torch >= 2.5
+    from torch.utils.data.dataloader import _T_co as T_co
 
 from modalities.dataloader.samplers import ResumableBatchSampler
 

From da5862eca9d1cd410c9354cf374609cf3b57835d Mon Sep 17 00:00:00 2001
From: Felix Stollenwerk <felix.stollenwerk@ai.se>
Date: Tue, 12 Nov 2024 11:11:20 +0100
Subject: [PATCH 2/2] fix(github actions): pin torch version to 2.4

---
 .github/workflows/tests_full.yml | 2 +-
 README.md                        | 4 ++--
 pyproject.toml                   | 2 +-
 3 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/.github/workflows/tests_full.yml b/.github/workflows/tests_full.yml
index e62120a4..7215313a 100644
--- a/.github/workflows/tests_full.yml
+++ b/.github/workflows/tests_full.yml
@@ -23,7 +23,7 @@ jobs:
         sudo apt-get update
         sudo apt-get install curl -y       # required by coveralls
         sudo apt-get install git -y
-        python -m pip install torch
+        python -m pip install torch~=2.4.1
         python -m pip install --upgrade pip setuptools wheel
         export FLASH_ATTENTION_SKIP_CUDA_BUILD=TRUE
         python -m pip install -e .[tests]
diff --git a/README.md b/README.md
index 7d415210..ec5961e0 100644
--- a/README.md
+++ b/README.md
@@ -59,7 +59,7 @@ Currently, the flash attention dependency cannot be installed without torch bein
 Until the flash attention developers fix this, we have to run
 
 ```sh
-pip install torch
+pip install torch~=2.4.1
 ```
 beforehand.
 
@@ -75,7 +75,7 @@ pip install -e .
 To install Modalities via pip, run
 
 ```sh
-pip install torch
+pip install torch~=2.4.1
 pip install modalities
 ```
 
diff --git a/pyproject.toml b/pyproject.toml
index cb9d91e8..b6531dd4 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -6,7 +6,7 @@ description = "Modalities, a PyTorch-native framework for distributed and reprod
 readme = "README.md"
 dependencies = [
     "numpy<2.0",
-    "torch>=2.3",
+    "torch~=2.4.1",
     "packaging",
     "tqdm",
     "pyyaml",