</span><span class="o">.</span><span class="n">size</span><span class="p">()[</span>
                    <span class="mi">0</span><span class="p">])</span><span class="o">.</span><span class="n">detach</span><span class="p">()</span><span class="o">.</span><span class="n">cpu</span><span class="p">()</span><span class="o">.</span><span class="n">numpy</span><span class="p">()</span>

        <span class="c1"># Storing the predictions</span>
        <span class="bp">self</span><span class="o">.</span><span c