Oh okay cool - will definitely check out that WGAN example - looks super helpful.
Last question (I promise!) - the GAN has a 2-component loss, real-vs-fake, for which I’m using the normal nn.BCELoss() , and an image style (class) classification loss , called style_criterion - I’m using a nn.MultiLabelSoftMarginLoss() [as the goal is for the discriminator to view each class as equiprobable]. I’m also using the retain_graph=True argument in all my backward() calls (for both of the losses) except the last one.
Does the below seem right?
style_labels = torch.FloatTensor(self.batch_size)
style_labels = style_labels.cuda()
#batch_styles = a Tensor of shape batch_sizex1
# contains the image style (class) labels of the current batch
style_labels = style_labels.copy_(batch_styles)
style_labels = Variable(style_labels.copy_(batch_styles))

gen_style_labels = torch.FloatTensor(self.batch_size)
# Dummy label as generator not trained on class labels
if self.cuda:
    gen_style_labels = gen_style_labels.cuda()
gen_style_labels = Variable(gen_style_labels)
gen_style_labels = gen_style_labels.fill_(real_label) 


# Discriminator
err_disc_style = style_criterion(output_styles, style_labels)
err_disc_style.backward(retain_graph=True) 
# later
disc_err = err_disc_real + err_disc_fake + err_disc_style

####

# Generator
err_gen_style = style_criterion(batch_labels,gen_style_labels)
err_gen_style.backward() # is the last backward() so don't need retain_graph=True
gen_err +=  err_gen_style # gen_err originally the conventional generator loss


I wrote this in line with the Tensorflow implementation of this GAN, but when I managed to (mostly) run it, it seemed to really mess with the generator as its output images stayed exactly the same across epochs (as opposed to the DCGAN version (with just the regular nn.BCELoss()), whose images improved from epoch to epoch.
Any ideas? Have I made a stupid mistake or something? Unfortunately the Tensorflow implementation of this GAN is confusing and I’m not a huge Tensorflow fan, so would like to implement it in the far nicer Pytorch 
Thanks again, you are a real lifesaver!