Created by: sidgoyal78

This PR modifies the code in Abhinav's PR https://github.com/PaddlePaddle/Paddle/pull/10472 to follow the new API