Update Gerador.py

Gerador.py

@@ -32,16 +32,22 @@ class Gerador(nn.Module):
             nn.Tanh() # Valor entre -1 e 1
         )
 
-    def forward(self, x):
-        x = self.encoder(x)
+    def forward(self, x, return_embedding=False, only_embedding=False):
+        embedding = self.encoder(x)
 
+        if only_embedding:
+            return embedding
+        
         # Ruído gaussiano, mesmo tamanho espacial de x
-        ruido = torch.randn(x.size(0), 1, x.size(2), x.size(3), device=x.device)
+        ruido = torch.randn(embedding.size(0), 1, embedding.size(2), embedding.size(3), device=embedding.device)
         ruido_feat = self.ruido_proj(ruido)
 
         # Concatena ruído aos features da imagem
-        x = torch.cat([x, ruido_feat], dim=1)
+        x = torch.cat([embedding, ruido_feat], dim=1)
 
         x = self.super(x)
 
-        return x
+        if return_embedding:
+            return (embedding, x)
+        else:
+            return x