text: import torch
  import torch.nn as nn
  import torch.nn.functional as F
  from torch.optim import AdamW
  
  from utils.loss import ContentLoss, AdversialLoss
  from utils.transforms import get_default_transforms
  from utils.datasets import get_dataloader
  from models.discriminator import Discriminator
  from models.generator import Generator
  
  
  def train():
      device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
  
      # Config
      batch_size = 9
      image_size = 256
      learning_rate = 1e-3
      beta1, beta2 = (.5, .99)
      weight_decay = 1e-3
      epochs = 10
  
      # Models
      netD = Discriminator().to(device)
      netG = Generator().to(device)
  
      optimizerD = AdamW(netD.parameters(), lr=learning_rate, betas=(beta1, beta2), weight_decay=weight_decay)
      optimizerG = AdamW(netG.parameters(), lr=learning_rate, betas=(beta1, beta2), weight_decay=weight_decay)
  
      # Labels
      cartoon_labels = torch.ones (batch_size, 1, image_size // 4, image_size // 4).to(device)
      fake_labels    = torch.zeros(batch_size, 1, image_size // 4, image_size // 4).to(device)
  
      # Loss functions
      content_loss = ContentLoss(device)
      adv_loss     = AdversialLoss(cartoon_labels, fake_labels)
      BCE_loss     = nn.BCELoss().to(device)
  
      # Dataloaders
      real_dataloader    = get_dataloader("./datasets/real_images",           size = image_size, bs = batch_size)
      cartoon_dataloader = get_dataloader("./datasets/cartoon_images",        size = image_size, bs = batch_size)
      edge_dataloader    = get_dataloader("./datasets/cartoon_images_smooth", size = image_size, bs = batch_size)
  
      # --------------------------------------------------------------------------------------------- #
      # Training Loop
  
      # Lists to keep track of progress
      img_list = []
      G_losses = []
      D_losses = []
      iters = 0
  
      tracked_images = next(iter(real_dataloader))[0].to(device)
  
      print("Starting Training Loop...")
      # For each epoch.
      for epoch in range(epochs):
          # For each batch in the dataloader.
          for i, ((cartoon_data, _), (edge_data, _), (real_data, _)) in enumerate(zip(cartoon_dataloader, edge_dataloader, real_dataloader )):
  
              ############################
              # (1) Update D network: maximize log(D(x)) + log(1 - D(G(z)))
              ###########################
              
              # Reset Discriminator gradient.
              netD.zero_grad()
  
              # Format batch.
              cartoon_data   = cartoon_data.to(device)
              edge_data      = edge_data.to(device)
              real_data      = real_data.to(device)
  
              # Generate image
              generated_data = netG(real_data)
  
              # Forward pass all batches through D.
              cartoon_pred   = netD(cartoon_data)      #.view(-1)
              edge_pred      = netD(edge_data)         #.view(-1)
              generated_pred = netD(generated_data)    #.view(-1)
  
              print(generated_data.is_cuda, real_data.is_cuda)
  
              # Calculate discriminator loss on all batches.
              errD = adv_loss(cartoon_pred, generated_pred, edge_pred)
              
              # Calculate gradients for D in backward pass
              errD.backward()
              D_x = cartoon_pred.mean().item() # Should be close to 1
  
              # Update D
              optimizerD.step()
  
  
              ############################
              # (2) Update G network: maximize log(D(G(z)))
              ###########################
              
              # Reset Generator gradient.
              netG.zero_grad()
              
              # Since we just updated D, perform another forward pass of all-fake batch through D
              generated_pred = netD(generated_data) #.view(-1)
  
              # Calculate G's loss based on this output
              print(generated_data.is_cuda, real_data.is_cuda)
              print("generated_pred:", generated_pred.is_cuda, "cartoon_labels:", cartoon_labels.is_cuda)
              errG = BCE_loss(generated_pred, cartoon_labels) + content_loss(generated_data, real_data)
  
              # Calculate gradients for G
              errG.backward()
  
              D_G_z2 = generated_pred.mean().item() # Should be close to 1
              
              # Update G
              optimizerG.step()
              
              # ---------------------------------------------------------------------------------------- #
  
              # Output training stats
              if i % 50 == 0:
                  print('[%d/%d][%d/%d]\tLoss_D: %.4f\tLoss_G: %.4f\tD(x): %.4f\tD(G(z)): %.4f / %.4f'
                      % (epoch, epochs, i, len(real_dataloader),
                          errD.item(), errG.item(), D_x, D_G_z1, D_G_z2))
  
              # Save Losses for plotting later
              G_losses.append(errG.item())
              D_losses.append(errD.item())
  
              # Check how the generator is doing by saving G's output on tracked_images
              if (iters % 500 == 0) or ((epoch == epochs-1) and (i == len(dataloader)-1)):
                  with torch.no_grad():
                      fake = netG(tracked_images).detach().cpu()
                  img_list.append(vutils.make_grid(fake, padding=2, normalize=True))
  
              iters += 1
  
  
  
  if __name__ == "__main__":
      train()
  