From a782368bbdbce34ca6c59c0b7ff59a1f86de98e1 Mon Sep 17 00:00:00 2001 From: Jingya HUANG <44135271+JingyaHuang@users.noreply.github.com> Date: Tue, 30 Apr 2024 15:38:34 +0200 Subject: [PATCH] Update optimum/neuron/modeling_diffusion.py Co-authored-by: Wenchen Li <9028430+neo@users.noreply.github.com> --- optimum/neuron/modeling_diffusion.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/optimum/neuron/modeling_diffusion.py b/optimum/neuron/modeling_diffusion.py index ca488c47a..1f74e49a3 100644 --- a/optimum/neuron/modeling_diffusion.py +++ b/optimum/neuron/modeling_diffusion.py @@ -853,9 +853,15 @@ def __init__( def forward( self, input_ids: torch.Tensor, - *args, - **kwargs, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, ): + if attention_mask is not None: + assert torch.equal(torch.ones_like(attention_mask), attention_mask), "attention_mask is expected to be only all ones" + if output_hidden_states is not None: + assert bool(self.text_encoder.config.output_hidden_states) == bool(output_hidden_states), "output_hidden_states is expected to be consistent with how it was compiled" + input_ids = input_ids.to(torch.long) # dummy generator uses long int for tracing inputs = (input_ids,)