IndexError using nonzero as an index

Bidski: <p>I am trying to reduce a tensor down to only the non-zero elements. If I am reading the docs correctly I should be able to do the following</p>
<pre><code class="lang-auto"># y_true has torch.Size([1, 6, 480, 640])
idx = torch.nonzero(torch.any(torch.greater(y_true, 0.0), dim=1), as_tuple=True)
y_true = y_true[idx] 
</code></pre>
<p>However, when I do, I get this error</p>
<pre><code class="lang-auto">IndexError: index 324 is out of bounds for dimension 1 with size 6
</code></pre>
<p>What have I missed?</p>

KFrank: <p>Hi Alex!</p>
<aside class="quote no-group quote-modified" data-username="Bidski" data-post="1" data-topic="129388" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/b/ee7513/48.png" class="avatar"> Bidski:</div>
<blockquote>
<pre><code class="lang-auto">idx = torch.nonzero(torch.any(torch.greater(y_true, 0.0), dim=1), as_tuple=True)
</code></pre>
<p>What have I missed?</p>
</blockquote>
</aside>
<p>The problem is that <code>torch.any (t, dim = 1)</code> “reduces” away a<br>
dimension, that is, it returns a result with one fewer dimension than<br>
has <code>t</code>.</p>
<p>In your case, <code>torch.any()</code> returns a tensor with shape <code>[1, 480, 640]</code>.<br>
<code>torch.nonzero()</code> then produces <code>idx</code>, a tuple of length 3, where <code>idx[1]</code><br>
corresponds to the size-480 dimension of <code>y_true</code>, and therefore can<br>
contain values up to 479.  But you then use <code>idx[1]</code> to index into the<br>
size-6 dimension of <code>y_true</code>, hence the <code>IndexError</code>.</p>
<p>I’m not sure what your use case is, but let me note that we know that<br>
<code>idx[0]</code> contains only values that are equal to 0 (because the 0th<br>
dimension of <code>y_true</code> has size 1).  So your code would “work,” for<br>
example, with</p>
<pre data-code-wrap="python"><code class="lang-python">y_true = y_true[(idx[0], idx[0], idx[1], idx[2])]
</code></pre>
<p>(but this probably isn’t what you want).</p>
<p>(Note, there is nothing subtle going on by using <code>idx[0]</code> twice – it’s just<br>
a hack where we use <code>idx[0]</code> as a handy all-zero tensor that can validly<br>
index into the size-6 dimension of <code>y_true</code>.)</p>
<p>If you could illustrate your desired use case with a complete, runnable<br>
script (including hard-coded or random data and dimensions much<br>
smaller than 480 and 640) that uses loops, if necessary, to produce<br>
the result you want, we can probably help you figure out how to do<br>
it with no-loop tensor operations.</p>
<p>Best.</p>
<p>K. Frank</p>

Bidski: <p>You are correct, replicating <code>idx[0]</code> is not desirable in this case. The first dimension having size 1 is just happenstance (this is the batch dimension and will ultimately be larger than 1).</p>
<p>I am basically trying to replicate this tensorflow code</p>
<pre><code class="lang-auto">idx = tf.squeeze(tf.where(tf.reduce_any(tf.greater(y_true, 0.0), axis=-1)), axis=-1)
y_true = tf.gather(y_true, idx)
</code></pre>