CUDA memory leak

Hi, I am trying to train several models in parallel using torch 's Since my setup has multiple GPUs, I pass a device also to my training task and the model is trained on that particular device.

The problem I face is RuntimeError: CUDA error: out of memory
after a while.

This happens after several models are trained and I can clearly see using watch nvidia-smi that the GPU memory accumulates over time.

I have posted a minimal example below which also leads to the stated issue.

Is there something needs to be done at the end of run() to clean up the memory?

Pytorch version: torch 1.6.0

import torch
from torch.multiprocessing import Pool
import random

class SimpleModule(torch.nn.Module):
    def __init__(self, input_: int, output: int):
        self.linear = torch.nn.Linear(input_, output)

    def forward(self, input_: torch.Tensor):
        return self.linear(input_)

def run(device: str):
    model = SimpleModule(10, 10).to(device)
    optimizer = torch.optim.SGD(lr=1e+2,
    loss_criterion = torch.nn.MSELoss()
    for i in range(int(1e+4)):
        inputs = torch.rand(5, 10).to(device)
        outputs = torch.rand(5, 10).to(device)
        preds = model.forward(inputs)
        loss = loss_criterion(preds, outputs)

if __name__ == "__main__":
    torch.multiprocessing.set_start_method("spawn", force=True)
    n_tasks = 500
    available_device = ["cpu"] if not torch.cuda.is_available() else [f"cuda:{i}" for i in
    devices = random.choices(available_device, k=n_tasks)
    pool = Pool(processes=50), devices)