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/3] 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/3] 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/3] 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),
     },