YannAhlgrim/wilds-ijepa
1
1
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

gradient checkpointing

Browse Source
This commit is contained in:
Yann Ahlgrim
2026-06-18 16:37:14 +02:00
parent 772acd8f0d
commit 918c1597bc
1 changed files with 2 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/train_supervised.py
+2
View File
@@ -214,6 +214,7 @@ def main(args, resume_preempt=False):
v_args = args["validation"] v_args = args["validation"]
es_args = o_args["early_stopping"] es_args = o_args["early_stopping"]
accum_steps = o_args.get("gradient_accumulation_steps", 1) 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") folder = resolve_log_dir(args, stage="train")
tag = l_args["write_tag"] tag = l_args["write_tag"]
@@ -245,6 +246,7 @@ def main(args, resume_preempt=False):
patch_size=mk_args["patch_size"], patch_size=mk_args["patch_size"],
crop_size=d_args["crop_size"], crop_size=d_args["crop_size"],
model_name=m_args["model_name"], model_name=m_args["model_name"],
use_gradient_checkpointing=use_gradient_checkpointing,
) )
representation_type = m_args.get("representation_type", "last_avgpool") representation_type = m_args.get("representation_type", "last_avgpool")