Training Of Localisation Network and Sampler

[tarun-munjal]

How are the weights getting trained in localisation network and sampler.
In original paper by { jaderberg, simonyan, zisserman, korayk } of Spatial Transformer Network , they have mentioned that the sampler and localisation networks weights can be trained using sub gradients. But in you code i haven't find any of the back propagation mechanism on both of these. Initially random weights are assigned and all the work is being done by the CNN you have implemented. I ran this code for 500 epochs on rotated a data set of MNIST but the results were very disappointing.
Can you please tell me where is that training done of Sampler, Grid generator and Localisation Network ?\