[tests] Fix failing gradient tests