Update discriminatorModel.py

discriminatorModel.py

@@ -1,21 +1,6 @@
-
 import torch
 import torch.nn as nn
 
-class Embedding(nn.Module):
-    def __init__(self, size_in, size_out):
-        super(Embedding, self).__init__()
-        self.text_embedding = nn.Sequential(
-            nn.Linear(size_in, size_out),
-            nn.BatchNorm1d(1),
-            nn.LeakyReLU(0.2, inplace=True)
-        )
-
-    def forward(self, x, text):
-        embed_out = self.text_embedding(text)
-        embed_out_resize = embed_out.repeat(4, 1, 4, 1).permute(1, 3, 0, 2)  # Resize to match the discriminator input size
-        out = torch.cat([x, embed_out_resize], 1)  # Concatenate text embedding with the input feature map
-        return out
 
 # The Discriminator model
 class Discriminator(nn.Module):
@@ -37,7 +22,22 @@ class Discriminator(nn.Module):
         self.output = nn.Sequential(
             nn.Conv2d(512 + self.embed_out_dim, 1, 4, 1, 0, bias=False), nn.Sigmoid()
         )
-
+        
+    class Embedding(nn.Module):
+        def __init__(self, size_in, size_out):
+            super(Embedding, self).__init__()
+            self.text_embedding = nn.Sequential(
+                nn.Linear(size_in, size_out),
+                nn.BatchNorm1d(1),
+                nn.LeakyReLU(0.2, inplace=True)
+            )
+    
+        def forward(self, x, text):
+            embed_out = self.text_embedding(text)
+            embed_out_resize = embed_out.repeat(4, 1, 4, 1).permute(1, 3, 0, 2)  # Resize to match the discriminator input size
+            out = torch.cat([x, embed_out_resize], 1)  # Concatenate text embedding with the input feature map
+            return out
+        
     def _create_layer(self, size_in, size_out, kernel_size=4, stride=2, padding=1, normalize=True):
         layers = [nn.Conv2d(size_in, size_out, kernel_size=kernel_size, stride=stride, padding=padding)]
         if normalize: