Hi there,
My framework contains two parts: 1). a selection mechanism 2). a predictor. The input data is a sequence but not all items in the sequence are useful. So I just use TopK items in the sequence to predict something.
But I failed to calculate the gradients since using top K indices block gradients flow to the selection model. (second part of the code)
A toy code as following:
seq = torch.randn(10, 15, 8)                  # batch * seq_length * feature
sel_model = nn.Linear(8, 1)                   # the selection mechanism
scores = sel_model(seq).squeeze() 

# using topK indices bring the autograd problem
val, ind = torch.topk(scores, 3, dim=1, largest=True)
ind_exp = ind.unsqueeze(-1).expand(10, 3, 8)
selected_seq = torch.gather(seq, 1, ind_exp) 

feature = selected_seq.mean(-1)
pred_model = nn.Linear(3, 1)
pred = pred_model(feature).squeeze()
loss = pred.sum()  # assume there is a loss
loss.backward()

print('sel_model grad：', sel_model.weight.grad)  # grad is None
print('select scores grad：', scores.grad)        # grad is None
print('pred_model grad：', pred_model.weight.grad)

So how to allow the gradients back to the selection model?