for epoch in range(EPOCH):

        for i, data in enumerate(train_loader, 0):

            net.train()

            model_out = net(
                Variable(data["left_aug"]), Variable(data["right_aug"])
            )
            REC_loss, disp_smooth_loss, lr_loss = get_loss(model_out)
      
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

image811×251 20.2 KB