commit 723193ec16897e8389a2ff7cf916a4a7e1ec564a
author Eddie Yan <eddiey@nvidia.com> Wed Sep 28 22:30:42 2022 +0000
committer PyTorch MergeBot <pytorchmergebot@users.noreply.github.com> Wed Sep 28 22:30:44 2022 +0000
tree 6fcd84122f219fa4ffd2dda29fa9b2cb7d5f7ae5
parent 01add6e2884347409a27efa16dcbf5b355ec4bd5

[cuDNN][cuDNN v8 API] Fix 3d convolution_add_relu in V8 (#85055)

Fix for issue uncovered in #84948

CC @ngimel
Pull Request resolved: https://github.com/pytorch/pytorch/pull/85055
Approved by: https://github.com/ngimel

aten/src/ATen/native/cudnn/Conv_v8.cpp

diff --git a/aten/src/ATen/native/cudnn/Conv_v8.cpp b/aten/src/ATen/native/cudnn/Conv_v8.cpp
index 2ad8d4f..8ecf5a2 100644
--- a/aten/src/ATen/native/cudnn/Conv_v8.cpp
+++ b/aten/src/ATen/native/cudnn/Conv_v8.cpp
@@ -654,7 +654,7 @@
     bool allow_tf32) {
   if (output.numel() == 0) { return; }
   if (at::native::cudnnv8_enabled_check_debug()) {
-    auto bias_ = bias.view({1, bias.numel(), 1, 1});
+    auto bias_ = input.ndimension() == 4 ? bias.view({1, bias.numel(), 1, 1}) : bias.view({1, bias.numel(), 1, 1, 1});
     run_fused_conv(input, output, weight, z, bias_,
       alpha, stride, padding, dilation,
       groups, benchmark, deterministic, allow_tf32);