Stop non-finite steps before optimizer updates (#4226)