From 7b1f2da828302ff938af76311ecf68a4ce0d7c2e Mon Sep 17 00:00:00 2001
From: Laurent Mazare <laurent.mazare@gmail.com>
Date: Wed, 6 Sep 2023 18:39:39 +0200
Subject: [PATCH] Cudnn fix. (#758)

---
 candle-core/src/cudnn.rs | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/candle-core/src/cudnn.rs b/candle-core/src/cudnn.rs
index 235ad6e3..dd466ba2 100644
--- a/candle-core/src/cudnn.rs
+++ b/candle-core/src/cudnn.rs
@@ -54,8 +54,8 @@ pub(crate) fn launch_conv2d<
     let x_shape = [
         params.b_size as i32,
         params.c_in as i32,
-        params.i_w as i32,
         params.i_h as i32,
+        params.i_w as i32,
     ];
     // Note that `src` already starts at the proper offset.
     let x = if src_l.is_contiguous() {
@@ -75,14 +75,14 @@ pub(crate) fn launch_conv2d<
         [
             params.c_out as i32,
             params.c_in as i32,
-            params.k_w as i32,
             params.k_h as i32,
+            params.k_w as i32,
         ],
     )?;
     let (w_out, h_out) = (params.out_w() as i32, params.out_h() as i32);
     let y = cudnn.create_4d_tensor(
         cudarc::cudnn::sys::cudnnTensorFormat_t::CUDNN_TENSOR_NCHW,
-        [params.b_size as i32, params.c_out as i32, w_out, h_out],
+        [params.b_size as i32, params.c_out as i32, h_out, w_out],
     )?;
     let conv2d = Conv2dForward {
         conv: &conv,