Recently, I found a piece of code written in Tensorflow. It serves to overriding tf.round()'s gradient with tf.identity()'s using gradient_override_map function:

import tensorflow as tf
G = tf.get_default_graph()
def quantize(x, k):
n = float(2**k - 1)
with G.gradient_override_map({“Round”: “Identity”}):
return tf.round(x * n) / n

And I want to implement it in pytorch but have no idea…Should I define a new operation?? And could anyone provide some advice on that?