Upload model

modeling_t5mimo.py

@@ -1420,6 +1420,7 @@ class T5MIMOForConditionalGeneration(T5PreTrainedModel):
         output_attentions: Optional[bool] = None,
         output_hidden_states: Optional[bool] = None,
         return_dict: Optional[bool] = None,
+        use_conv: Optional[bool] = True,
     ) -> Union[Tuple[torch.FloatTensor], Seq2SeqLMOutput]:
         r"""
         labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
@@ -1517,6 +1518,9 @@ class T5MIMOForConditionalGeneration(T5PreTrainedModel):
 
         sequence_output = decoder_outputs[0]
 
+        if use_conv:
+            sequence_output = self.conv_block(sequence_output)
+            
         # Set device for model parallelism
         if self.model_parallel:
             torch.cuda.set_device(self.encoder.first_device)
@@ -1528,7 +1532,6 @@ class T5MIMOForConditionalGeneration(T5PreTrainedModel):
             # See https://github.com/tensorflow/mesh/blob/fa19d69eafc9a482aff0b59ddd96b025c0cb207d/mesh_tensorflow/transformer/transformer.py#L586
             sequence_output = sequence_output * (self.model_dim**-0.5)
 
-        sequence_output = self.conv_block(sequence_output)
         lm_logits = self.lm_head(sequence_output)
 
         loss = None