# CUDA Running out of memory after a few batches in an epoch

**URL:** <https://discuss.pytorch.org/t/cuda-running-out-of-memory-after-a-few-batches-in-an-epoch/143914>\
**Category:** nlp\
**Created:** [February 11, 2022, 1:19pm UTC](https://discuss.pytorch.org/t/cuda-running-out-of-memory-after-a-few-batches-in-an-epoch/143914 "2022-02-11T13:19:13Z")\
**Posts on this page:** 2\
**Page:** 1

<div class="post-metadata">

**Author:** ![DebadityaPal](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/debadityapal/32/46579_2.png) [@DebadityaPal](https://discuss.pytorch.org/u/DebadityaPal)\
**Post date:** [February 11, 2022, 1:19pm UTC](https://discuss.pytorch.org/t/cuda-running-out-of-memory-after-a-few-batches-in-an-epoch/143914/1 "2022-02-11T13:19:13Z")

</div>

This is my training function:

```auto
def train():

  device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
  model.to(device)
  model.train()
  optim = torch.optim.AdamW(model.parameters(), lr=5e-5)

  for epoch in range(10):
      with tqdm(dataloader, unit=" batch", leave=True, position=0) as tepoch:
        for i, data in enumerate(tepoch):
            inputs = tokenizer(data['text'], padding=True, truncation=True, return_tensors='pt')
            lm_labels = inputs['input_ids'].to(device)
            sarcasm_labels = data['label'].to(device)
            masked_tokens = mask_inputs(inputs['input_ids']).to(device)
            attention_mask = inputs['attention_mask'].to(device)
            
            outputs = model(masked_tokens, attention_mask=attention_mask, labels=lm_labels, next_sentence_label=sarcasm_labels)

            optim.zero_grad()
            # extract loss
            loss = outputs.loss
            # calculate loss for every parameter that needs grad update
            loss.backward()
            # update parameters
            optim.step()
            # print relevant info to progress bar
            tepoch.set_description(f'Epoch {epoch}')
            tepoch.set_postfix(loss=loss.detach().item())
            torch.cuda.empty_cache()

```

The training begins fine, however after a few batches in the first epoch itself, I get the Runtime Error of Cuda being out of memory.  
The model is a Huggingface BertForPretraining Model.  
Can someone help me with this?

Edit: For some reason, if I reduce the size of my Dataset, the problem is solved, but I cannot figure out why this is the case.

---

<div class="post-metadata">

**Author:** ![ptrblck](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ptrblck/32/1823_2.png) [@ptrblck](https://discuss.pytorch.org/u/ptrblck)\
**Post date:** [February 11, 2022, 4:58pm UTC](https://discuss.pytorch.org/t/cuda-running-out-of-memory-after-a-few-batches-in-an-epoch/143914/2 "2022-02-11T16:58:59Z")

</div>

> [@DebadityaPal](#):
>
> For some reason, if I reduce the size of my Dataset, the problem is solved, but I cannot figure out why this is the case.

This might point to a memory increase in each iteration, which might not be causing the OOM anymore, if you are reducing the number of iterations.  
Check the memory usage in your code e.g. via `torch.cuda.memory_summary()` or `torch.cuda.memory_allocated()` inside the training iterations and try to narrow down where the increase happens (you should also see that e.g. `loss.backward()` reduces the memory usage).