<h2 dir="auto"><g-emoji class="g-emoji" alias="bug" fallback-src="https://github.githubassets.com/images/icons/emoji/unicode/1f41b.png">🐛</g-emoji> Bug</h2>
<p dir="auto">When indexing into a tensor with another tensor, around 10 times the size of the array is used in memory on the GPU.</p>
<h2 dir="auto">To Reproduce</h2>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="import torch


def torch_memory(device):
    # Checks and prints GPU memory
    print(f'{torch.cuda.memory_allocated(device)/1024/1024:.2f} MB USED')
    print(f'{torch.cuda.memory_reserved(device)/1024/1024:.2f} MB RESERVED')
    print(f'{torch.cuda.max_memory_allocated(device)/1024/1024:.2f} MB USED MAX')
    print(f'{torch.cuda.max_memory_reserved(device)/1024/1024:.2f} MB RESERVED MAX')
    print('')


device = torch.device(0)

a = torch.randn((1, 32, 24, 512, 512), dtype=torch.float32, device=device)  # 768 MB tensor
torch_memory(device)

indices = a &lt; 0  # 192 MB tensor
torch_memory(device)

a[indices] = 0.  # Allocates 10 times of 768 MB
torch_memory(device)"><pre class="notranslate"><span class="pl-k">import</span> <span class="pl-s1">torch</span>


<span class="pl-k">def</span> <span class="pl-en">torch_memory</span>(<span class="pl-s1">device</span>):
    <span class="pl-c"># Checks and prints GPU memory</span>
    <span class="pl-en">print</span>(<span class="pl-s">f'<span class="pl-s1"><span class="pl-kos">{</span><span class="pl-s1">torch</span>.<span class="pl-s1">cuda</span>.<span class="pl-en">memory_allocated</span>(<span class="pl-s1">device</span>)<span class="pl-c1">/</span><span class="pl-c1">1024</span><span class="pl-c1">/</span><span class="pl-c1">1024</span>:.2f<span class="pl-kos">}</span></span> MB USED'</span>)
    <span class="pl-en">print</span>(<span class="pl-s">f'<span class="pl-s1"><span class="pl-kos">{</span><span class="pl-s1">torch</span>.<span class="pl-s1">cuda</span>.<span class="pl-en">memory_reserved</span>(<span class="pl-s1">device</span>)<span class="pl-c1">/</span><span class="pl-c1">1024</span><span class="pl-c1">/</span><span class="pl-c1">1024</span>:.2f<span class="pl-kos">}</span></span> MB RESERVED'</span>)
    <span class="pl-en">print</span>(<span class="pl-s">f'<span class="pl-s1"><span class="pl-kos">{</span><span class="pl-s1">torch</span>.<span class="pl-s1">cuda</span>.<span class="pl-en">max_memory_allocated</span>(<span class="pl-s1">device</span>)<span class="pl-c1">/</span><span class="pl-c1">1024</span><span class="pl-c1">/</span><span class="pl-c1">1024</span>:.2f<span class="pl-kos">}</span></span> MB USED MAX'</span>)
    <span class="pl-en">print</span>(<span class="pl-s">f'<span class="pl-s1"><span class="pl-kos">{</span><span class="pl-s1">torch</span>.<span class="pl-s1">cuda</span>.<span class="pl-en">max_memory_reserved</span>(<span class="pl-s1">device</span>)<span class="pl-c1">/</span><span class="pl-c1">1024</span><span class="pl-c1">/</span><span class="pl-c1">1024</span>:.2f<span class="pl-kos">}</span></span> MB RESERVED MAX'</span>)
    <span class="pl-en">print</span>(<span class="pl-s">''</span>)


<span class="pl-s1">device</span> <span class="pl-c1">=</span> <span class="pl-s1">torch</span>.<span class="pl-en">device</span>(<span class="pl-c1">0</span>)

<span class="pl-s1">a</span> <span class="pl-c1">=</span> <span class="pl-s1">torch</span>.<span class="pl-en">randn</span>((<span class="pl-c1">1</span>, <span class="pl-c1">32</span>, <span class="pl-c1">24</span>, <span class="pl-c1">512</span>, <span class="pl-c1">512</span>), <span class="pl-s1">dtype</span><span class="pl-c1">=</span><span class="pl-s1">torch</span>.<span class="pl-s1">float32</span>, <span class="pl-s1">device</span><span class="pl-c1">=</span><span class="pl-s1">device</span>)  <span class="pl-c"># 768 MB tensor</span>
<span class="pl-en">torch_memory</span>(<span class="pl-s1">device</span>)

<span class="pl-s1">indices</span> <span class="pl-c1">=</span> <span class="pl-s1">a</span> <span class="pl-c1">&lt;</span> <span class="pl-c1">0</span>  <span class="pl-c"># 192 MB tensor</span>
<span class="pl-en">torch_memory</span>(<span class="pl-s1">device</span>)

<span class="pl-s1">a</span>[<span class="pl-s1">indices</span>] <span class="pl-c1">=</span> <span class="pl-c1">0.</span>  <span class="pl-c"># Allocates 10 times of 768 MB</span>
<span class="pl-en">torch_memory</span>(<span class="pl-s1">device</span>)</pre></div>
<p dir="auto">This is the output that I get.</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="768.00 MB USED
768.00 MB RESERVED
768.00 MB USED MAX
768.00 MB RESERVED MAX

960.00 MB USED
960.00 MB RESERVED
960.00 MB USED MAX
960.00 MB RESERVED MAX

960.00 MB USED
8662.00 MB RESERVED
8643.00 MB USED MAX
8662.00 MB RESERVED MAX"><pre class="notranslate"><code class="notranslate">768.00 MB USED
768.00 MB RESERVED
768.00 MB USED MAX
768.00 MB RESERVED MAX

960.00 MB USED
960.00 MB RESERVED
960.00 MB USED MAX
960.00 MB RESERVED MAX

960.00 MB USED
8662.00 MB RESERVED
8643.00 MB USED MAX
8662.00 MB RESERVED MAX
</code></pre></div>
<p dir="auto">The same thing happens when I use (supposed to be equivalent)</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="a.index_put_((indices,), torch.tensor([0.]).to(device))"><pre class="notranslate"><span class="pl-s1">a</span>.<span class="pl-en">index_put_</span>((<span class="pl-s1">indices</span>,), <span class="pl-s1">torch</span>.<span class="pl-en">tensor</span>([<span class="pl-c1">0.</span>]).<span class="pl-en">to</span>(<span class="pl-s1">device</span>))</pre></div>
<p dir="auto">and even when I just index into the array without assigning new values to it, a lot of memory is used.</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="a[indices]"><pre class="notranslate"><span class="pl-s1">a</span>[<span class="pl-s1">indices</span>]</pre></div>
<h2 dir="auto">Expected behavior</h2>
<p dir="auto">I'm not sure how much memory indexing/assignment like this is supposed to use, but surely not 10 times.</p>
<p dir="auto">I tried this with CuPy arrays and they use 2-3 times the size of the array in memory.</p>
<h2 dir="auto">Environment</h2>
<ul dir="auto">
<li>PyTorch Version (e.g., 1.0): 1.1, 1.4, 1.5</li>
<li>OS (e.g., Linux): Linux</li>
<li>How you installed PyTorch (<code class="notranslate">conda</code>, <code class="notranslate">pip</code>, source): conda</li>
<li>Build command you used (if compiling from source): N/A</li>
<li>Python version: 3.7</li>
<li>CUDA/cuDNN version: 10.1, 7.6</li>
<li>GPU models and configuration: V100</li>
<li>Any other relevant information: N/A</li>
</ul>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="In [1]: x = torch.randn(20000, 20000,device='cuda')

In [2]: torch.cuda.max_memory_allocated()
Out[2]: 1600126976

In [3]: m = ~torch.eye(20000, dtype=torch.bool, device='cuda')

In [4]: torch.cuda.max_memory_allocated()
Out[4]: 2401239040

In [5]: x[m]
Out[5]:
tensor([ 0.6415,  0.2951, -0.7630,  ...,  1.0940,  0.3103,  0.3631],
       device='cuda:0')

In [6]: torch.cuda.max_memory_allocated()
Out[6]: 10001317888

In [7]: torch.__version__
Out[7]: '1.4.0a0+91c6d2e'"><pre class="notranslate"><span class="pl-v">In</span> [<span class="pl-c1">1</span>]: <span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-s1">torch</span>.<span class="pl-en">randn</span>(<span class="pl-c1">20000</span>, <span class="pl-c1">20000</span>,<span class="pl-s1">device</span><span class="pl-c1">=</span><span class="pl-s">'cuda'</span>)

<span class="pl-v">In</span> [<span class="pl-c1">2</span>]: <span class="pl-s1">torch</span>.<span class="pl-s1">cuda</span>.<span class="pl-en">max_memory_allocated</span>()
<span class="pl-v">Out</span>[<span class="pl-c1">2</span>]: <span class="pl-c1">1600126976</span>

<span class="pl-v">In</span> [<span class="pl-c1">3</span>]: <span class="pl-s1">m</span> <span class="pl-c1">=</span> <span class="pl-c1">~</span><span class="pl-s1">torch</span>.<span class="pl-en">eye</span>(<span class="pl-c1">20000</span>, <span class="pl-s1">dtype</span><span class="pl-c1">=</span><span class="pl-s1">torch</span>.<span class="pl-s1">bool</span>, <span class="pl-s1">device</span><span class="pl-c1">=</span><span class="pl-s">'cuda'</span>)

<span class="pl-v">In</span> [<span class="pl-c1">4</span>]: <span class="pl-s1">torch</span>.<span class="pl-s1">cuda</span>.<span class="pl-en">max_memory_allocated</span>()
<span class="pl-v">Out</span>[<span class="pl-c1">4</span>]: <span class="pl-c1">2401239040</span>

<span class="pl-v">In</span> [<span class="pl-c1">5</span>]: <span class="pl-s1">x</span>[<span class="pl-s1">m</span>]
<span class="pl-v">Out</span>[<span class="pl-c1">5</span>]:
<span class="pl-en">tensor</span>([ <span class="pl-c1">0.6415</span>,  <span class="pl-c1">0.2951</span>, <span class="pl-c1">-</span><span class="pl-c1">0.7630</span>,  ...,  <span class="pl-c1">1.0940</span>,  <span class="pl-c1">0.3103</span>,  <span class="pl-c1">0.3631</span>],
       <span class="pl-s1">device</span><span class="pl-c1">=</span><span class="pl-s">'cuda:0'</span>)

<span class="pl-v">In</span> [<span class="pl-c1">6</span>]: <span class="pl-s1">torch</span>.<span class="pl-s1">cuda</span>.<span class="pl-en">max_memory_allocated</span>()
<span class="pl-v">Out</span>[<span class="pl-c1">6</span>]: <span class="pl-c1">10001317888</span>

<span class="pl-v">In</span> [<span class="pl-c1">7</span>]: <span class="pl-s1">torch</span>.<span class="pl-s1">__version__</span>
<span class="pl-v">Out</span>[<span class="pl-c1">7</span>]: <span class="pl-s">'1.4.0a0+91c6d2e'</span></pre></div>
<p dir="auto">Titan XP<br>
CUDA: 10.1</p>
<p dir="auto">Why does it allocate about 4.75x of the indexed tensor, rather than just 1x?</p>
<p dir="auto">cc <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/ngimel/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/ngimel">@ngimel</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/VitalyFedyunin/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/VitalyFedyunin">@VitalyFedyunin</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/mruberry/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/mruberry">@mruberry</a></p>