title: Understanding why memory allocation occurs during inference, backpropagation, and model update
body: <p>In the process of tracking down a GPU OOM error, I made the following checkpoints in my Pytorch code (running on Google Colab P100):</p>
  <pre><code>learning_rate = 0.001
  num_epochs = 50
  
  device = torch.device(&quot;cuda:0&quot; if torch.cuda.is_available() else &quot;cpu&quot;)
  
  print('check 1')
  !nvidia-smi | grep MiB | awk '{print $9 $10 $11}'
  
  model = MyModel()
  
  print('check 2')
  !nvidia-smi | grep MiB | awk '{print $9 $10 $11}'
  
  model = model.to(device)
  
  print('check 3')
  !nvidia-smi | grep MiB | awk '{print $9 $10 $11}'
  
  optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
  
  print('check 4')
  !nvidia-smi | grep MiB | awk '{print $9 $10 $11}'
  
  for epoch in range(num_epochs):
      train_running_loss = 0.0
      train_accuracy = 0.0
  
      model = model.train()
  
      print('check 5')
      !nvidia-smi | grep MiB | awk '{print $9 $10 $11}'
  
      ## training step
      for i, (name, output_array, input) in enumerate(trainloader):
          
          output_array = output_array.to(device)
          input = input.to(device)
          comb = torch.zeros(1,1,100,1632).to(device)
  
          print('check 6')
          !nvidia-smi | grep MiB | awk '{print $9 $10 $11}'
  
          ## forward + backprop + loss
          output = model(input, comb)
  
          print('check 7')
          !nvidia-smi | grep MiB | awk '{print $9 $10 $11}'
  
          loss = my_loss(output, output_array)
  
          print('check 8')
          !nvidia-smi | grep MiB | awk '{print $9 $10 $11}'
  
          optimizer.zero_grad()
  
          print('check 9')
          !nvidia-smi | grep MiB | awk '{print $9 $10 $11}'
  
          loss.backward()
  
          print('check 10')
          !nvidia-smi | grep MiB | awk '{print $9 $10 $11}'
  
          ## update model params
          optimizer.step()
  
          print('check 11')
          !nvidia-smi | grep MiB | awk '{print $9 $10 $11}'
  
          train_running_loss += loss.detach().item()
  
          print('check 12')
          !nvidia-smi | grep MiB | awk '{print $9 $10 $11}'
  
          temp = get_accuracy(output, output_array)
  
          print('check 13')
          !nvidia-smi | grep MiB | awk '{print $9 $10 $11}'
  
          train_accuracy += temp     
  </code></pre>
  <p>with the following output:</p>
  <pre><code>check 1
  2MiB/16160MiB
  check 2
  2MiB/16160MiB
  check 3
  3769MiB/16160MiB
  check 4
  3769MiB/16160MiB
  check 5
  3769MiB/16160MiB
  check 6
  3847MiB/16160MiB
  check 7
  6725MiB/16160MiB
  check 8
  6725MiB/16160MiB
  check 9
  6725MiB/16160MiB
  check 10
  9761MiB/16160MiB
  check 11
  16053MiB/16160MiB
  check 12
  16053MiB/16160MiB
  check 13
  16053MiB/16160MiB
  check 6
  16053MiB/16160MiB
  check 7
  16071MiB/16160MiB
  check 8
  16071MiB/16160MiB
  check 9
  16071MiB/16160MiB
  check 10
  16071MiB/16160MiB
  ---------------------------------------------------------------------------
  RuntimeError                              Traceback (most recent call last)
  &lt;ipython-input-11-f566d09448f9&gt; in &lt;module&gt;()
       65 
       66         ## update model params
  ---&gt; 67         optimizer.step()
       68 
       69         print('check 11')
  
  3 frames
  /usr/local/lib/python3.7/dist-packages/torch/optim/optimizer.py in wrapper(*args, **kwargs)
       86                 profile_name = &quot;Optimizer.step#{}.step&quot;.format(obj.__class__.__name__)
       87                 with torch.autograd.profiler.record_function(profile_name):
  ---&gt; 88                     return func(*args, **kwargs)
       89             return wrapper
       90 
  
  /usr/local/lib/python3.7/dist-packages/torch/autograd/grad_mode.py in decorate_context(*args, **kwargs)
       26         def decorate_context(*args, **kwargs):
       27             with self.__class__():
  ---&gt; 28                 return func(*args, **kwargs)
       29         return cast(F, decorate_context)
       30 
  
  /usr/local/lib/python3.7/dist-packages/torch/optim/adam.py in step(self, closure)
      116                    lr=group['lr'],
      117                    weight_decay=group['weight_decay'],
  --&gt; 118                    eps=group['eps'])
      119         return loss
  
  /usr/local/lib/python3.7/dist-packages/torch/optim/_functional.py in adam(params, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, state_steps, amsgrad, beta1, beta2, lr, weight_decay, eps)
       92             denom = (max_exp_avg_sqs[i].sqrt() / math.sqrt(bias_correction2)).add_(eps)
       93         else:
  ---&gt; 94             denom = (exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_(eps)
       95 
       96         step_size = lr / bias_correction1
  
  RuntimeError: CUDA out of memory. Tried to allocate 2.32 GiB (GPU 0; 15.78 GiB total capacity; 11.91 GiB already allocated; 182.75 MiB free; 14.26 GiB reserved in total by PyTorch)
  </code></pre>
  <p>It makes sense to me that <code>model = model.to(device)</code> creates 3.7G of memory.</p>
  <p>But why does running the model <code>output = model(input, comb)</code> create another 3G of memory?
  <br>And then <code>loss.backward()</code> creates another 3G of memory?
  <br>And then <code>optimizer.step()</code> creates another 6.3G of memory?</p>
  <p>I would appreciate it if someone could explain how the PyTorch GPU memory allocation model is working in this example.</p>