Stochastic Weight Averaging for steps instead of epochs

[MaveriQ]

Hi

I am looking for guidance on how to adapt the SWA callback for a step based averaging and scheduler update. In its present form, both actions take place in the function on_train_epoch_start. Some of this will have to be moved into 'training_step' function.

I will appreciate any pointers on how to go about that. Thanks