I tried the following after noticing this post - What's the difference between a[0][1] and a[0 , 1] and it worked for me.
output_variable[batch_idx, word_idx, 0, matching_idx] = cosine
So, I am just curious to know why pytorch interprets x[i, j] and x[i][j] differently? As soumith mentioned in that post,

When you index with x[i][j], then an intermediate Tensor x[i] is created first, and the operation [j] is applied on it. If you index with x[i, j] then there’s no intermediate operation.

I want to know when we should use x[i][j] and when x[i, j]?