Thanks, something like this?
for epoch in range(epochs):
    for batch in dataloader:
        start_scale = scaler.get_scale()
        # do stuff
        end_scale = scaler.get_scale()
        if not end_scale < start_scale:
            scheduler.step()

I’m not quite sure if I follow where to place the checks. If I put the start_scale check at the beginning of each batch loop and the end_scale check at the very end, for nearly every batch I get:
start_scale == 65536.0
end_scale == 32.0
And so it would never get to the scheduler.step()
Edit: Actually by the start of epoch 2 it looks like it settles so that both start_scale and end_scale == 32.0
Should I similarly suppress the swa_scheduler step or is that left alone?