diff --git a/.github/workflows/test_and_deploy.yml b/.github/workflows/test_and_deploy.yml
index 75c252d..f8e29f0 100644
--- a/.github/workflows/test_and_deploy.yml
+++ b/.github/workflows/test_and_deploy.yml
@@ -30,6 +30,8 @@ jobs:
     needs: [linting, manifest]
     name: ${{ matrix.os }} py${{ matrix.python-version }}
     runs-on: ${{ matrix.os }}
+    env:
+      KERAS_BACKEND: torch
     strategy:
       fail-fast: false
       matrix:
diff --git a/brainglobe_workflows/brainmapper/main.py b/brainglobe_workflows/brainmapper/main.py
index 66e6ef6..89b600f 100644
--- a/brainglobe_workflows/brainmapper/main.py
+++ b/brainglobe_workflows/brainmapper/main.py
@@ -154,7 +154,6 @@ def run_all(args, what_to_run, atlas):
             args.model_weights,
             args.install_path,
             args.model,
-            args.n_free_cpus,
         )
         if what_to_run.classify:
             if points is None: