Merge pull request #31 from Geopipe/multigpu-disableloss-fix

More robust disabling of disabled losses.

1 files changed, 1 insertions, 1 deletions

diff --git a/train.py b/train.py
index 231662c..368dc9f 100755
--- a/train.py
+++ b/train.py
@@ -61,7 +61,7 @@ for epoch in range(start_epoch, opt.niter + opt.niter_decay + 1):
 
         # calculate final loss scalar
         loss_D = (loss_dict['D_fake'] + loss_dict['D_real']) * 0.5
-        loss_G = loss_dict['G_GAN'] + loss_dict['G_GAN_Feat'] + loss_dict['G_VGG']
+        loss_G = loss_dict['G_GAN'] + loss_dict.get('G_GAN_Feat',0) + loss_dict.get('G_VGG',0)
 
         ############### Backward Pass ####################
         # update generator weights