Dear eqy,
Thank you for your reply.
Here is my training loop :
# transfer the model to the gpu
model.to(device)

#define optimizer
optimizer = torch.optim.Adam(model.parameters(), lr=0.000001)

# define loss function
distortion = nn.MSELoss().cuda()
        
beta = 0.00001
n_epochs = 400

for epoch in range(0, n_epochs):
    running_loss = 0.0

    for i_batch, data in enumerate(dataloader):
        batch_images = data[0].to(device).float()
        batch_masks = data[1].to(device).float()
        [decoded_images, x_quantized] = model(batch_images, batch_mask, 1, True)
        optimizer.zero_grad()
        loss_dist = distortion(decoded_images, batch_images)
        loss_bit = entropy_dist(x_quantized, model.phi, model.var)
        loss = beta * loss_dist + loss_bit
        loss.backward()
        optimizer.step()
        running_loss += loss.item()

    running_loss = running_loss/len(dataloader)

I don’t know if I should prevent batch_mask from beeing stored for backpropagation (batch_mask.detach()). batch_mask does not appear in the loss but is an entry to my custom convolutional layers, and is used to modify the main input batch_images.