Fix Sampling inside gradient loop issue

[bekleyis95]

This MR is an extension on MR 147, where the algorithm is fixed to support sampling inside the gradient step loop instead of outside. The changes are made for all the places uses gradient step loop which diverges from the methodology outlined in the papers.