Here is my training procedure:
def train(train_gen, model, criterion, optimizer, epoch, contrast_w=0):
    epoch_loss = 0
    for iteration, batch in enumerate(train_gen, 1):
        torch.cuda.empty_cache()
        #input, targetdeblur, targetsr
        LR_Blur = batch[0]
        LR_Deblur = batch[1]
        HR = batch[2]
        LR_Blur = LR_Blur.to(device)
        LR_Deblur = LR_Deblur.to(device)
        HR = HR.to(device)

        if opt.isTest == True:
            test_Tensor = torch.cuda.FloatTensor().resize_(1).zero_()+1
        else:
            test_Tensor = torch.cuda.FloatTensor().resize_(1).zero_()
        if opt.gated == True:
            gated_Tensor = torch.cuda.FloatTensor().resize_(1).zero_()+1
        else:
            gated_Tensor = torch.cuda.FloatTensor().resize_(1).zero_()

        print("Before", iteration, LR_Blur.shape, HR.shape, LR_Deblur.shape)
        [lr_deblur, sr] = model(LR_Blur, gated_Tensor, test_Tensor)
        print("after", lr_deblur.shape, sr.shape, LR_Blur.shape, gated_Tensor.shape, test_Tensor.shape)

        loss1 = criterion(lr_deblur, LR_Deblur)
        loss2 = criterion(sr, HR)
        mse = loss2 + opt.lambda_db * loss1

I print the shape of the data before and after entering the model.
Also, I print the shape at the beginning and the end in the forward function of Net.
def forward(self, x, gated, isTest):
        print("IN model", x.shape, gated.shape)
        
        ...... ...... ...... 

        print("[return]", deblur_out.shape, recon_out.shape)
        return deblur_out, recon_out

Here are the outputs:
Before 1 torch.Size([16, 3, 24, 24]) torch.Size([16, 3, 96, 96]) torch.Size([16, 3, 24, 24])
IN model torch.Size([8, 3, 24, 24]) torch.Size([1])
[return] torch.Size([8, 3, 24, 24]) torch.Size([8, 3, 96, 96])
after torch.Size([8, 3, 24, 24]) torch.Size([8, 3, 96, 96]) torch.Size([16, 3, 24, 24]) torch.Size([1]) torch.Size([1])