From a5567827cd5cb0d756a32ed76c487792367bf823 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Bob=20Dr=C3=B6ge?= <b.e.droge@rug.nl>
Date: Thu, 12 Mar 2026 16:45:19 +0100
Subject: [PATCH 1/4] set PyTorch CPU limits for
 aarch64/{generic,neoverse_n1,neoverse_v1}

---
 eb_hooks.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/eb_hooks.py b/eb_hooks.py
index 82673010..066a5b82 100644
--- a/eb_hooks.py
+++ b/eb_hooks.py
@@ -2021,6 +2021,11 @@ def set_maximum(parallel, max_value):
         '*': (divide_by_factor, 2),
         CPU_TARGET_A64FX: (set_maximum, 1),
     },
+    'PyTorch': {
+        CPU_TARGET_AARCH64_GENERIC: (divide_by_factor, 2),
+        CPU_TARGET_NEOVERSE_N1: (divide_by_factor, 2),
+        CPU_TARGET_NEOVERSE_V1: (divide_by_factor, 2),
+    },
     'QuantumESPRESSO': {
         CPU_TARGET_A64FX: (set_maximum, 6),
     },

From cc9e9e84703f08d30448ff4300b893aa88990b59 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Bob=20Dr=C3=B6ge?= <b.e.droge@rug.nl>
Date: Thu, 12 Mar 2026 23:12:30 +0100
Subject: [PATCH 2/4] use eb_hooks.py from git repo

---
 EESSI-install-software.sh | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/EESSI-install-software.sh b/EESSI-install-software.sh
index e4834545..ccd96f76 100755
--- a/EESSI-install-software.sh
+++ b/EESSI-install-software.sh
@@ -396,6 +396,8 @@ else
             # load EasyBuild module (will be installed if it's not available yet)
             source ${TOPDIR}/load_easybuild_module.sh ${eb_version}
 
+            export EASYBUILD_HOOKS=$TOPDIR/eb_hooks.py
+
             ${EB} --show-config
 
             echo_green "All set, let's start installing some software with EasyBuild v${eb_version} in ${EASYBUILD_INSTALLPATH}..."

From fa76fc6d29e72c2dfd9b9e909e4d0602c304ca15 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Bob=20Dr=C3=B6ge?= <b.e.droge@rug.nl>
Date: Sun, 15 Mar 2026 08:45:27 +0100
Subject: [PATCH 3/4] use 8 cores for PyTorch on A64FX

---
 eb_hooks.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/eb_hooks.py b/eb_hooks.py
index 066a5b82..c7791bd6 100644
--- a/eb_hooks.py
+++ b/eb_hooks.py
@@ -2023,6 +2023,7 @@ def set_maximum(parallel, max_value):
     },
     'PyTorch': {
         CPU_TARGET_AARCH64_GENERIC: (divide_by_factor, 2),
+        CPU_TARGET_A64FX: (set_maximum, 8),
         CPU_TARGET_NEOVERSE_N1: (divide_by_factor, 2),
         CPU_TARGET_NEOVERSE_V1: (divide_by_factor, 2),
     },

From ae3ae48b9469b9c27fd08ecbebc25a2adc352c3c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Bob=20Dr=C3=B6ge?= <b.e.droge@rug.nl>
Date: Mon, 22 Jun 2026 15:13:45 +0200
Subject: [PATCH 4/4] allow 41 failing tests for PyTorch 2.9.1 on A64FX

---
 eb_hooks.py | 14 +++++++++-----
 1 file changed, 9 insertions(+), 5 deletions(-)

diff --git a/eb_hooks.py b/eb_hooks.py
index 813f196e..32072c5c 100644
--- a/eb_hooks.py
+++ b/eb_hooks.py
@@ -1711,11 +1711,15 @@ def pre_test_hook_increase_max_failed_tests_arm_PyTorch(self, *args, **kwargs):
     See https://github.com/EESSI/software-layer/issues/461
     """
     cpu_target = get_eessi_envvar('EESSI_SOFTWARE_SUBDIR')
-    if self.name == 'PyTorch' and self.version == '2.1.2':
-        if get_cpu_architecture() == AARCH64:
-            self.cfg['max_failed_tests'] = 10
-        if cpu_target in [CPU_TARGET_CASCADELAKE, CPU_TARGET_ICELAKE, CPU_TARGET_SAPPHIRE_RAPIDS]:
-            self.cfg['max_failed_tests'] = 4
+    if self.name == 'PyTorch':
+        if self.version == '2.1.2':
+            if get_cpu_architecture() == AARCH64:
+                self.cfg['max_failed_tests'] = 10
+            if cpu_target in [CPU_TARGET_CASCADELAKE, CPU_TARGET_ICELAKE, CPU_TARGET_SAPPHIRE_RAPIDS]:
+                self.cfg['max_failed_tests'] = 4
+        if self.version == '2.9.1':
+            if cpu_target == CPU_TARGET_A64FX:
+                self.cfg['max_failed_tests'] = 41
 
 
 def pre_test_hook_ignore_failing_tests_OpenBabel_a64fx(self, *args, **kwargs):