diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index a320b98..efb9461 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -26,7 +26,7 @@ jobs:
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip
-        python -m pip install torch==2.3.0 --index-url https://download.pytorch.org/whl/cpu
+        python -m pip install torch==2.3.1 torchvision==0.18.1 --index-url https://download.pytorch.org/whl/cpu
         python -m pip install scipy
     - name: Test with pytest
       run: |
diff --git a/meshgpt_pytorch/data.py b/meshgpt_pytorch/data.py
index 94b01b7..243fbdf 100644
--- a/meshgpt_pytorch/data.py
+++ b/meshgpt_pytorch/data.py
@@ -13,8 +13,8 @@
 
 from einops import rearrange, reduce
 
-from beartype import beartype
 from beartype.typing import Tuple, List, Callable, Dict
+from meshgpt_pytorch.typing import typecheck
 
 from torchtyping import TensorType
 
@@ -38,7 +38,7 @@ def identity(t):
 # you would decorate your Dataset class with this
 # and then change your `data_kwargs = ["text_embeds", "vertices", "faces"]`
 
-@beartype
+@typecheck
 def cache_text_embeds_for_dataset(
     embed_texts_fn: Callable[[List[str]], Tensor],
     max_text_len: int,
@@ -157,7 +157,7 @@ def __getitem__(self, idx):
 # you would decorate your Dataset class with this function
 # and then change your `data_kwargs = ["vertices", "faces", "face_edges"]`
 
-@beartype
+@typecheck
 def cache_face_edges_for_dataset(
     max_edges_len: int,
     cache_path: str = './face_edges_cache',
@@ -261,7 +261,7 @@ def __getitem__(self, idx):
 # dataset
 
 class DatasetFromTransforms(Dataset):
-    @beartype
+    @typecheck
     def __init__(
         self,
         folder: str,
diff --git a/meshgpt_pytorch/trainer.py b/meshgpt_pytorch/trainer.py
index bd76d25..2f7c3fa 100644
--- a/meshgpt_pytorch/trainer.py
+++ b/meshgpt_pytorch/trainer.py
@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from pathlib import Path
 from functools import partial
 from packaging import version
diff --git a/meshgpt_pytorch/version.py b/meshgpt_pytorch/version.py
index 96e3ce8..8e3c933 100644
--- a/meshgpt_pytorch/version.py
+++ b/meshgpt_pytorch/version.py
@@ -1 +1 @@
-__version__ = '1.4.0'
+__version__ = '1.4.1'
diff --git a/setup.py b/setup.py
index dd45769..f8af194 100644
--- a/setup.py
+++ b/setup.py
@@ -23,10 +23,10 @@
     'accelerate>=0.25.0',
     'beartype',
     "huggingface_hub>=0.21.4",
-    'classifier-free-guidance-pytorch>=0.6.2',
-    'einops>=0.7.0',
-    'einx[torch]>=0.1.3',
-    'ema-pytorch>=0.5.0',
+    'classifier-free-guidance-pytorch>=0.6.7',
+    'einops>=0.8.0',
+    'einx[torch]>=0.3.0',
+    'ema-pytorch>=0.5.1',
     'environs',
     'gateloop-transformer>=0.2.2',
     'jaxtyping',