Merge pull request #266 from Modalities/fix/tests

fix: towards torch 2.5 compatibility / fix unit tests (github actions)

.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]

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
 ```
 

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",

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