Since you are dealing with invalid gradients as NaNs, I would check the code for operations which might cause Infs or NaNs.
E.g. could you check, if you could possibly divide by zero or a small number in:
new_locs[:, i, ...] = 2 * (new_locs[:, i, ...] / (shape[len(shape)-i-1] - 1) - 0.5)

which could blow up the activations and might cause problems?