Updated ganbert.py

ganbert.py

@@ -58,8 +58,6 @@ class GAN(PreTrainedModel):
       self.model_name = self.all_checkpoints[config.model_number]
       self.parent_config = AutoConfig.from_pretrained(self.model_name)
       self.hidden_size = int(self.parent_config.hidden_size)
-      self.ns = config.noise_size
-      self.dv = config.device
       # Define the number and width of hidden layers
       self.hidden_levels_g = [self.hidden_size for i in range(0, config.num_hidden_layers_g)]
       self.hidden_levels_d = [self.hidden_size for i in range(0, config.num_hidden_layers_d)]
@@ -73,7 +71,7 @@ class GAN(PreTrainedModel):
       # Put everything in the GPU if available
       # print(self.generator,self.discriminator)
       self.transformer = AutoModel.from_pretrained(self.model_name,output_attentions=True)
-      
+      self.config = config
       if config.device == 'cuda':    
         self.generator.cuda()
         self.discriminator.cuda()
@@ -81,7 +79,8 @@ class GAN(PreTrainedModel):
   def forward(self,**kwargs):
     # Encode real data in the Transformer
     # real_batch_size = input_ids.shape[0]
-    model_outputs = self.transformer(**kwargs)
+    model_outputs = self.transformer(output_hidden_states = self.config.output_hidden_states,\
+                                     output_attentions = self.config.output_attentions,**kwargs)
     # print('got transformer output')
     # hidden_states = torch.mean(model_outputs[0],dim=1)
     # noise = torch.zeros(real_batch_size, self.ns, device=self.dv).uniform_(0, 1).to(self.dv)