for mask, images in zip(roi_loader, data_loader):
images = images.to(device)
mask = mask.unsqueeze(1).to(device)
if is_add_noise:
images += 0.05 * torch.randn(images.size()).to(device)
# if final batch isn't equal to defined batch size in loader
batch_size = images.size()[0]
random = generate_noise(batch_size, **noise_hyperparams).to(device)
_, loss_d = trainer.train_step_discriminator(random, mask, images)
if train_dis:
trainer.backward_discriminator(loss_d)
random = generate_noise(batch_size, **noise_hyperparams).to(device)
gen_images, loss_g, info_loss = trainer.train_step_generator(
random, mask, images)
if train_gen:
trainer.backward_generator(loss_g)
'''
random = torch.randn(batch_size, noise_hyperparams['noise_dim']).to(device)
random_control = 2 * torch.rand(batch_size,
noise_hyperparams['cont_dim']
).to(device) - 1
random = torch.cat([random, random_control], dim=1)
for i in range(noise_hyperparams['n_disc']):
random_control = 2 * torch.rand(batch_size,
