title: Quickly find indices that have values larger than a threshold in Numpy/PyTorch
body: <p><strong>Task</strong></p>
  
  <p>Given a numpy or pytorch matrix, find the indices of cells that have values that are larger than a given threshold.</p>
  
  <p><strong>My implementation</strong></p>
  
  <pre><code>#abs_cosine is the matrix
  #sim_vec is the wanted
  
  sim_vec = []
  for m in range(abs_cosine.shape[0]):
      for n in range(abs_cosine.shape[1]):
          # exclude diagonal cells
          if m != n and abs_cosine[m][n] &gt;= threshold:
              sim_vec.append((m, n))
  </code></pre>
  
  <p><strong>Concerns</strong></p>
  
  <p><strong>Speed</strong>. All other computations are built on Pytorch, using <code>numpy</code> is already a compromise, because it has moved computations from GPU to CPU. Pure python <code>for</code> loops will make the whole process even worse (for small data set already 5 times slower). <strong>I was wondering if we can move the whole computation to Numpy (or pytorch) without invoking any <code>for</code> loops?</strong></p>
  
  <p>An improvement I can think of (but got stuck...)</p>
  
  <blockquote>
    <p>bool_cosine = abs_cosine > threshold</p>
  </blockquote>
  
  <p>which returns a boolean matrix of <code>True</code> and <code>False</code>. But I cannot find a way to quick retrieve the indices of the <code>True</code> cells.</p>