Hi~ I also encountered this problem. My code is like this:
# fd_prob: [batch_size, tgt_len, vocab_size], means words probability distribution
# bd_hyp: [batch_size, infer_len], means another output labels
# fd_bd_attn: [batch_size, tgt_len, infer_len], means edit probability distribution
# fd_p_gen: [batch_size, tgt_len, 1], means copy mode probability
batch_size, tgt_len, _ = fd_prob.size()
_, infer_len = bd_hyp.size()
# incorporate copy mode
for i in range(batch_size):
for j in range(tgt_len):
for k in range(infer_len):
fd_prob[i][j][bd_hyp[i][k]] += (1 - fd_p_gen[i][j][0]) * fd_bd_attn[i][j][k]
loss = criterion(fd_prob, ground_truth)
loss.backward()
I modified output distribution by a 3-layer loop which contained lots of index, and I found that forward cost less than 1s while backward cost more than 1.5min, which is unacceptable. I think it’s due to lots of index. Have you found the solution or is there an elegant way to do this?
Also see Indexing is very slow for backpropagation