Make the context contiguous if required.

huggingface · Aug 29, 2023 · 5d2df98 · 5d2df98
1 parent 383d1be
commit 5d2df98
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/candle-examples/examples/stable-diffusion/attention.rs b/candle-examples/examples/stable-diffusion/attention.rs
@@ -208,9 +208,9 @@ impl CrossAttention {
  fn forward(&self, xs: &Tensor, context: Option<&Tensor>) -> Result<Tensor> {
  let _enter = self.span.enter();
  let query = self.to_q.forward(xs)?;
- let context = context.unwrap_or(xs);
- let key = self.to_k.forward(context)?;
- let value = self.to_v.forward(context)?;
+ let context = context.unwrap_or(xs).contiguous()?;
+ let key = self.to_k.forward(&context)?;
+ let value = self.to_v.forward(&context)?;
  let query = self.reshape_heads_to_batch_dim(&query)?;
  let key = self.reshape_heads_to_batch_dim(&key)?;
  let value = self.reshape_heads_to_batch_dim(&value)?;