classifier on generator for tga (#3716)

facebookresearch · Jun 15, 2021 · e93a00f · e93a00f
1 parent fe36d88
commit e93a00f
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/projects/style_gen/modules.py b/projects/style_gen/modules.py
@@ -160,6 +160,10 @@ def forward(self, input, encoder_state, embedded_input=None, incr_state=None):
                 )
             )
         tensor = tensor + self.position_embeddings(positions).expand_as(tensor)
+
+        if self.variant == 'bart':
+            tensor = self.norm_embeddings(tensor)
+
         tensor = self.dropout(tensor)  # --dropout
 
         new_incr_state = {}