From dcf708559da9f5d9f07a1394200236aec261e820 Mon Sep 17 00:00:00 2001 From: Laurent Mazare Date: Wed, 6 Sep 2023 08:49:28 +0200 Subject: [PATCH] Fix for cudnn to work with img2img. (#753) --- candle-examples/examples/stable-diffusion/unet_2d_blocks.rs | 2 ++ 1 file changed, 2 insertions(+) diff --git a/candle-examples/examples/stable-diffusion/unet_2d_blocks.rs b/candle-examples/examples/stable-diffusion/unet_2d_blocks.rs index 26a1035b..be258acb 100644 --- a/candle-examples/examples/stable-diffusion/unet_2d_blocks.rs +++ b/candle-examples/examples/stable-diffusion/unet_2d_blocks.rs @@ -754,6 +754,7 @@ impl UpBlock2D { let mut xs = xs.clone(); for (index, resnet) in self.resnets.iter().enumerate() { xs = Tensor::cat(&[&xs, &res_xs[res_xs.len() - index - 1]], 1)?; + xs = xs.contiguous()?; xs = resnet.forward(&xs, temb)?; } match &self.upsampler { @@ -855,6 +856,7 @@ impl CrossAttnUpBlock2D { let mut xs = xs.clone(); for (index, resnet) in self.upblock.resnets.iter().enumerate() { xs = Tensor::cat(&[&xs, &res_xs[res_xs.len() - index - 1]], 1)?; + xs = xs.contiguous()?; xs = resnet.forward(&xs, temb)?; xs = self.attentions[index].forward(&xs, encoder_hidden_states)?; }