SNIP should calculate gradients on indicators/mask, not on weights

[iamanigeeit]

Hello,

Thanks for implementing in PyTorch, but i believe there is a something wrong in the code.

In the original paper and implementation, loss is differentiated against the connection indicators and not the weights.

From Lee's original code in line 67:
grads = tf.gradients(loss, [mask_init[k] for k in prn_keys])

I understand you have weight = indicator * weight before computing gradients, but i can't see where you extract the gradients for the indicators only. I see you've posted on the pytorch forum about this but nobody has answered properly.