move out zero grad logic into separate function

move out zero grad logic into separate function #1284