diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 25db415..ea7fe5a 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -7,10 +7,12 @@ on:
 jobs:
   test:
     name: test
-    runs-on: ubuntu-latest
+    runs-on: ${{ matrix.os }}
     strategy:
       matrix:
-        python-version: ["3.9", "3.10", "3.11", "3.12", "3.13"]
+        os: [ubuntu-latest, ubuntu-24.04-arm, macos-latest]
+        # Only test older and newer supported versions
+        python-version: ["3.9", "3.13"]
       fail-fast: true
 
     steps:
@@ -37,7 +39,7 @@ jobs:
 
       - uses: codacy/codacy-coverage-reporter-action@89d6c85cfafaec52c72b6c5e8b2878d33104c699
         if: >-
-          ${{ matrix.python-version == '3.11' && (
+          ${{ matrix.python-version == '3.11' && matrix.os == 'ubuntu-latest' && (
           github.event_name != 'pull_request' ||
           github.event.pull_request.head.repo.full_name == github.repository) }}
         with:
diff --git a/pyproject.toml b/pyproject.toml
index 246299b..8912851 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [project]
 name = "deeprank-gnn-esm"
-version = "1.0.0"
+version = "1.0.1"
 description = "Graph Network for protein-protein interface including language model features"
 authors = [{ name = "Xiaotong Xu" }, { email = "x.xu1@uu.nl" }]
 maintainers = [{ name = "BonvinLab" }, { email = "bonvinlab.support@uu.nl" }]
diff --git a/src/deeprank_gnn/DataSet.py b/src/deeprank_gnn/DataSet.py
index 01c009b..ea602b8 100644
--- a/src/deeprank_gnn/DataSet.py
+++ b/src/deeprank_gnn/DataSet.py
@@ -11,6 +11,10 @@
 from .community_pooling import community_detection, community_pooling
 
 
+def _default_edge_feature_transform(x):
+    return np.tanh(-x / 2 + 2) + 1
+
+
 def DivideDataSet(dataset, percent=[0.8, 0.2], shuffle=True):
     """Divides the dataset into a training set and an evaluation set
 
@@ -101,7 +105,7 @@ def __init__(
         node_feature="all",
         edge_feature=["dist"],
         clustering_method="mcl",
-        edge_feature_transform=lambda x: np.tanh(-x / 2 + 2) + 1,
+        edge_feature_transform=_default_edge_feature_transform,
     ):
         """Class from which the hdf5 datasets are loaded.