[architecture] fixed dimension and cuda issues in the original DR-GAN architecture

bob/learn/pytorch/architectures/DRGANOriginal.py

@@ -322,12 +322,13 @@ class DRGANOriginal_discriminator(nn.Module):
       nn.BatchNorm2d(latent_dim),
       # ------------------------------------------
       
-      # --- average pool
       # input is (latent_dim)x6x6, output is latent_dimx1x1
       nn.AvgPool2d(6, stride=1),
+    )
 
-      # --- fully connected
-      #nn.Linear(320, (number_of_ids + conditional_dim + 1))
+    self.classifier =  nn.Sequential(
+      nn.Linear(self.latent_dim, (self.number_of_ids + self.conditional_dim + 1)),
+      nn.Sigmoid()
     )
     
 
@@ -339,17 +340,14 @@ class DRGANOriginal_discriminator(nn.Module):
 
     x: pyTorch Variable
       The minibatch of images to process.
-    """    
+    """
     if isinstance(x.data, torch.cuda.FloatTensor) and self.ngpu > 1:
       output_avgpool = nn.parallel.data_parallel(self.main, x, range(self.ngpu))
+      output_avgpool = output_avgpool.squeeze()
+      output = nn.parallel.data_parallel(self.classifier, output_avgpool, range(self.ngpu))
     else:
       output_avgpool = self.main(x)
-
-    input_linear = output_avgpool.squeeze()
-    classifier = nn.Sequential(
-                    nn.Linear(self.latent_dim, (self.number_of_ids + self.conditional_dim + 1)),
-                    nn.Sigmoid()
-                 )
-    output = classifier(input_linear)
+      output_avgpool = output_avgpool.squeeze()
+      output = self.classifier(output_avgpool)
 
     return output