Use multiple transformer layer in the same cross-attn blocks. (#653)

* Use multiple transformer layer in the same cross-attn blocks.

* Make the context contiguous if required.

1 files changed, 3 insertions, 3 deletions

diff --git a/candle-examples/examples/stable-diffusion/attention.rs b/candle-examples/examples/stable-diffusion/attention.rs
index 797542aa..1ae1bfc3 100644
--- 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)?;