diff --git a/src/diffusers/models/attention_processor.py b/src/diffusers/models/attention_processor.py index f88400da0333..a489814c4787 100644 --- a/src/diffusers/models/attention_processor.py +++ b/src/diffusers/models/attention_processor.py @@ -344,11 +344,14 @@ def get_attention_scores(self, query, key, attention_mask=None): beta=beta, alpha=self.scale, ) + del baddbmm_input if self.upcast_softmax: attention_scores = attention_scores.float() attention_probs = attention_scores.softmax(dim=-1) + del attention_scores + attention_probs = attention_probs.to(dtype) return attention_probs