diff --git a/config/process_base.config b/config/process_base.config
index bd0e3c14..598b0999 100644
--- a/config/process_base.config
+++ b/config/process_base.config
@@ -24,6 +24,9 @@ process{
   withLabel: mem_96 {
     memory = {check_memory(48.GB  + 48.GB * task.attempt, params.max_memory)}
   }
+  withLabel: mem_max {
+    memory = {task.attempt > 1 ? params.max_memory : check_memory(96.GB, params.max_memory)}
+  }
   withLabel: cpus_2  {
     cpus = {check_cpus(2, params.max_cpus)}
   }
diff --git a/modules/classify-celltypes.nf b/modules/classify-celltypes.nf
index 5cb155d8..67846a73 100644
--- a/modules/classify-celltypes.nf
+++ b/modules/classify-celltypes.nf
@@ -49,7 +49,7 @@ process classify_cellassign {
       mode: 'copy',
       pattern: "${cellassign_dir}"
     )
-  label 'mem_96'
+  label 'mem_max'
   label 'cpus_12'
   tag "${meta.library_id}"
   input: