drop tuple

Signed-off-by: Mayank Mishra <[email protected]>
IBM · Sep 11, 2024 · ca4a833 · ca4a833
1 parent 4967b9c
commit ca4a833
Showing 1 changed file with 0 additions and 1 deletion.
diff --git a/dolomite_engine/hf_models/mixins/dense_TP/main.py b/dolomite_engine/hf_models/mixins/dense_TP/main.py
@@ -60,7 +60,6 @@ def forward(
         cu_seqlens: torch.Tensor | None = None,
         max_seqlen: torch.Tensor | None = None,
     ) -> tuple | CausalLMOutputWithPast:
-        assert not output_attentions
         assert return_dict
 
         input_ids, position_ids, token_type_ids, labels, cu_seqlens, max_seqlen = self.prepare_inputs_for_model(