GPU memory consumption increases a lot at the first several iterations while training.

I am also facing this issue. I tried the solution



 apaszke:

If you add del loss, output at the end of the loop the memory usage will likely remain the same after the first iteration (what you see it probably a side-effect of Python’s scoping rules). It’s possible that cuDNN uses much less memory than the default backend.


But this did not have noticeable impact. My training loop is as follows:
def run(self):
    pb = tqdm(
        total=self._n,
        leave=self._keep,
        colour=self._bar_color,
        desc=self._desc
    )
    train_accuracies = 0.0
    train_ious = 0.0
    train_mses = 0.0
    for data_dict in self._generator:
        if self._eval_mode:
            self._model.eval()
        else:
            self._model.train()
        self._optimizer.zero_grad()
        pred = self._model(
            x_a=data_dict['x1'],
            x_b=data_dict['x2'],
        )
        im_acc = label_accuracy(
            prediction=pred['y1'],
            target=data_dict['y1']
        )
        im_iou = iou_score(
            predictions=pred['y1'],
            targets=data_dict['y1']
        )
        pc_mse = pc_mse_metric(
            prediction=pred['x2'],
            target=data_dict['x2']
        )
        train_accuracies += im_acc
        train_ious += im_iou
        train_mses += pc_mse
        if not self._eval_mode:
            loss = self._criterion(
                predictions=pred, targets={
                    'y1': data_dict['y1'],
                    'x2': data_dict['x2']
                }
            )

            loss.backward()
            self._optimizer.step()
        pb.update(n=1)
    k = "Validation" if self._eval_mode else "Training"
    self._metrics.append(
        {
            k: {
                'Accuracy': (train_accuracies / self._n),
                'mIoU': (train_ious / self._n),
                'MSE': (train_mses / self._n)
            }
        }
    )


Due this problem I am encountering RuntimeError: CUDA out of memory after a few iterations.
What am I doing wrong?