Thank you for your answer,
I am not sure of the following code behaviour :
def train(model, train_loader, criterion, optimizer, epoch):
    model.train()
    running_loss = 0
    for i,data in enumerate(train_loader):
        image = data['image']
        mask = data['mask']
        
        optimizer.zero_grad()
        output = model(image)

        loss = criterion(output, mask)
        loss.backward()

        optimizer.step()

(train is called every epoch)
I tought it was loading image an mask by pair and after computing the loss the pair was unloaded from GPU is that right ?
By the way my batch size is 2