@tom explains a nice way to achieve this in this post and also mentions the shortcomings of other approaches such as zeroing out the gradients of some elements.