# How to get round memory leak caused by accumulating loss in for loop?

**URL:** <https://discuss.pytorch.org/t/how-to-get-round-memory-leak-caused-by-accumulating-loss-in-for-loop/176224>\
**Category:** vision\
**Created:** [March 30, 2023, 3:16am UTC](https://discuss.pytorch.org/t/how-to-get-round-memory-leak-caused-by-accumulating-loss-in-for-loop/176224 "2023-03-30T03:16:07Z")\
**Posts on this page:** 2\
**Page:** 1

<div class="post-metadata">

**Author:** ![JimW](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/jimw/32/38000_2.png) [@JimW](https://discuss.pytorch.org/u/JimW)\
**Post date:** [March 30, 2023, 3:16am UTC](https://discuss.pytorch.org/t/how-to-get-round-memory-leak-caused-by-accumulating-loss-in-for-loop/176224/1 "2023-03-30T03:16:07Z")

</div>

Hi I have received a model from someone else， and I am trying to train it with our flow.

For this model, the most convenient way is to generate one label for a single sample. So I am doing something like this:

```auto
total_loss=0 
for i in range(len(dataset)):
     pred = model(dataset[i])
     gt = labels[i]
     if i%batch_size==0:
         total_loss.backward()
         optimizer.step()
         total_loss = 0
     else：
         total_loss += CrossEntropy(pred, gt) 

```

I understand it is not the best practice to accumulate loss in this way. My observation is : when I used batch\_size larger than 4, it will run for several epochs and then crash with CUDA Out of Memory error.

Is there a way in pytorch to get around this?  
I found some solutions online which is to add “.item()” to total\_loss, but I think it will make this node undifferentiable.

---

<div class="post-metadata">

**Author:** ![Joschka](https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/j/4af34b/32.png) [@Joschka](https://discuss.pytorch.org/u/Joschka)\
**Post date:** [March 30, 2023, 7:02am UTC](https://discuss.pytorch.org/t/how-to-get-round-memory-leak-caused-by-accumulating-loss-in-for-loop/176224/2 "2023-03-30T07:02:10Z")

</div>

I see two issues with your code.

1. You are never zeroing the gradients of your model in-between [optimization steps](https://pytorch.org/docs/stable/optim.html#taking-an-optimization-step). This is not necessarily related to your memory leak, but most likely still a problem.

2. I think if `i % batch_size` is equal to 0, the prediction made in that iteration and the correpsonding computation graph might never be cleared, because with your current control-flow, they are never subject to any backward pass. Computing and accumulating the loss in _every_ iteration should fix this.

Hope this helps:

```python
total_loss=0 
for i in range(len(dataset)):
     pred = model(dataset[i])
     gt = labels[i]
     total_loss += CrossEntropy(pred, gt)
     if i%batch_size==0:
         optimizer.zero_grad()
         total_loss.backward()
         optimizer.step()
         total_loss = 0

```