Hey thanks, it is helpful.
However, I’m still wondering what’s the correct way to use top-k values, as top-k will break the chain rule (not sure just in my opinion)? Could you provide some examples?
Below is a modified version of the previous code, using top-k values.
data = torch.randn(10, 15)  # batch * features
select_model = nn.Linear(15, 15)  # each feature has a score
scores = select_model(data)

# use top-3 features and mask the rest
val, ind = torch.topk(scores, 3, dim=1, largest=True)
masked_scores = torch.zeros_like(scores)
masked_scores.scatter_(1, ind, val)
masked_data = data * masked_scores  # mask some sequences

# I can use retain_grad to keep grad, but which one I should use?
# scores.retain_grad()
# masked_scores.retain_grad()
# masked_data.retain_grad()

# do the prediction task using the masked data
pred_model = nn.Linear(15, 1)
pred = pred_model(masked_data).squeeze()
loss = pred.sum()  # assume there is a loss
loss.backward()

print('select_model grad：', select_model.weight.grad.size())
print('select scores grad：', scores.grad)  # grad is None
print('masked scores grad：', masked_scores.grad)  # grad is None
print('pred_model grad：', pred_model.weight.grad.size())