From 918c1597bcbe2a50847af972a0035132a5119ebc Mon Sep 17 00:00:00 2001 From: Yann Ahlgrim Date: Thu, 18 Jun 2026 16:37:14 +0200 Subject: [PATCH] gradient checkpointing --- src/train_supervised.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/train_supervised.py b/src/train_supervised.py index 624fb90..cbe12b6 100644 --- a/src/train_supervised.py +++ b/src/train_supervised.py @@ -214,6 +214,7 @@ def main(args, resume_preempt=False): v_args = args["validation"] es_args = o_args["early_stopping"] accum_steps = o_args.get("gradient_accumulation_steps", 1) + use_gradient_checkpointing = m_args.get("use_gradient_checkpointing", True) folder = resolve_log_dir(args, stage="train") tag = l_args["write_tag"] @@ -245,6 +246,7 @@ def main(args, resume_preempt=False): patch_size=mk_args["patch_size"], crop_size=d_args["crop_size"], model_name=m_args["model_name"], + use_gradient_checkpointing=use_gradient_checkpointing, ) representation_type = m_args.get("representation_type", "last_avgpool")