I changed my code by defining two variables at the top of the for loop so that register_hook(), retain_grad() and backward() . will all be called on the same variable:
                    ranked_tensori = ranked_tensor[nonzeroidx.data[0][0]]
                    ranked_tensorj = ranked_tensor[idx]

This code was working until I refactored some unrelated code, now I am getting the following error:
raise RuntimeError(“cannot register a hook on a volatile variable”)