elif args.dataset == 'mpii':
    lsp_train_dataset = MPII('train')
    lsp_val_dataset = MPII('valid')  ## MPII('val') was present originally

# Dataset and the Dataloade
train_loader = torch.utils.data.DataLoader(lsp_train_dataset,
                                           batch_size=args.batch_size,
                                           shuffle=True)
val_save_loader = torch.utils.data.DataLoader(lsp_val_dataset,
                                              batch_size=args.val_batch_size)
val_eval_loader = torch.utils.data.DataLoader(lsp_val_dataset,
                                              batch_size=args.val_batch_size,
                                              shuffle=True)
#train_eval = torch.utils.data.DataLoader(lsp_train_dataset, batch_size=args.val_batch_size, shuffle=True)

pck = metrics.PCK(metrics.Options(256, 8))

# Loading on GPU, if available
if (args.use_gpu):
    generator_model = generator_model.to(fast_device)
    discriminator_model_conf = discriminator_model_conf.to(fast_device)
    discriminator_model_pose = discriminator_model_pose.to(fast_device)

# Cross entropy loss
#criterion = nn.CrossEntropyLoss()

# Setting the optimizer
if (args.optimizer_type == 'SGD'):
    optim_gen = optim.SGD(generator_model.parameters(),
                          lr=args.lr,
                          momentum=args.momentum)
# MPII
elif args.dataset == 'mpii':
    lsp_train_dataset = MPII('train')
    lsp_val_dataset = MPII('val')

# Dataset and the Dataloade
train_loader = torch.utils.data.DataLoader(lsp_train_dataset,
                                           batch_size=args.batch_size,
                                           shuffle=True)
val_save_loader = torch.utils.data.DataLoader(lsp_val_dataset,
                                              batch_size=args.val_batch_size)
val_eval_loader = torch.utils.data.DataLoader(lsp_val_dataset,
                                              batch_size=args.val_batch_size,
                                              shuffle=True)

pck = metrics.PCK(metrics.Options(256, config['generator']['num_stacks']))

# Loading on GPU, if available
if (args.use_gpu):
    generator_model = generator_model.to(fast_device)
    discriminator_model = discriminator_model.to(fast_device)

# Cross entropy loss
criterion = nn.CrossEntropyLoss()

# Setting the optimizer
if (args.optimizer_type == 'SGD'):
    optim_gen = optim.SGD(generator_model.parameters(),
                          lr=args.lr,
                          momentum=args.momentum)
    optim_disc = optim.SGD(discriminator_model.parameters(),