Optimizer is reset every epoch rather than during training stage switch

[JaydenTeoh]

In the paper, it is mentioned that "we also reset the optimizer state when training stages switch."

However, in

coconut/run.py

Lines 312 to 320 in d1c10b1

	if configs.reset_optimizer:
	del optimizer
	optimizer = optim.AdamW(
	parallel_model.parameters(),
	lr=configs.lr,
	weight_decay=configs.weight_decay,
	)

the optimizer is reset every epoch.

[Ber666]

Hi, sorry for the confusion. It's indeed reset at every epoch for simplicity. If you reset the optimizer at every stage, the performance should be similar (if not better).