Reversed to no_grad for head_init_scale

GuillaumeErhard · GuillaumeErhard · commit 2d09181a050b · 2025-07-04T01:44:35.000+02:00
diff --git a/train.py b/train.py
@@ -512,7 +512,7 @@ def main():
         **args.model_kwargs,
     )
     if args.head_init_scale is not None:
-        with torch.inference_mode():
+        with torch.no_grad():
             model.get_classifier().weight.mul_(args.head_init_scale)
             model.get_classifier().bias.mul_(args.head_init_scale)
     if args.head_init_bias is not None:

Original file line number	Diff line number	Diff line change
`@@ -512,7 +512,7 @@ def main():`
`512`	`512`	`**args.model_kwargs,`
`513`	`513`	`)`
`514`	`514`	`if args.head_init_scale is not None:`
`515`		`- with torch.inference_mode():`
	`515`	`+ with torch.no_grad():`
`516`	`516`	`model.get_classifier().weight.mul_(args.head_init_scale)`
`517`	`517`	`model.get_classifier().bias.mul_(args.head_init_scale)`
`518`	`518`	`if args.head_init_bias is not None:`