title: Pytorch: create a mask that is larger than the n-th quantile of each 2D tensor in a batch
body: <p>I have a <code>torch.Tensor</code> of shape <code>(2, 2, 2)</code> (can be bigger), where the values are normalized within range <code>[0, 1]</code>.</p>
  <p>Now I am given a positive integer <code>K</code>, which tells me that I need to create a mask where for each 2D tensor inside the batch, values are 1 if it is larger than <code>1/k</code> of all the values, and 0 elsewhere. The return mask also has shape <code>(2, 2, 2)</code>.</p>
  <p>For example, if I have a batch like this:</p>
  <pre class="lang-py prettyprint-override"><code>tensor([[[1., 3.],
           [2., 4.]],
          [[5., 7.],
           [9., 8.]]])
  </code></pre>
  <p>and let <code>K=2</code>, it means that I must mask the values where they are greater than 50% of all the values inside each 2D tensor.</p>
  <p>In the example, the 0.5 quantile is <code>2.5</code> and <code>7.5</code>, so this is the desired output:</p>
  <pre class="lang-py prettyprint-override"><code>tensor([[[0, 1],
           [0, 1]],
          [[0, 0],
           [1, 1]]])
  </code></pre>
  <p>I tried:</p>
  <pre class="lang-py prettyprint-override"><code>a = torch.tensor([[[0, 1],
                     [0, 1]],
                    [[0, 0],
                     [1, 1]]])
  quantile = torch.tensor([torch.quantile(x, 1/K) for x in a])
  torch.where(a &gt; val, 1, 0)
  </code></pre>
  <p>But this is the result:</p>
  <pre class="lang-py prettyprint-override"><code>tensor([[[0, 0],
           [0, 0]],
          [[1, 0],
           [1, 1]]])
  </code></pre>