Use FP32 compute type for FP16 convolutions (#1115)

OpenNMT · Mar 6, 2023 · 6e0f757 · 6e0f757
1 parent a21e4f4
commit 6e0f757
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/src/ops/conv1d_gpu.cu b/src/ops/conv1d_gpu.cu
@@ -50,7 +50,11 @@ namespace ctranslate2 {
                                                   /*stride_h=*/1, /*stride_w=*/_stride,
                                                   /*dilation_h=*/1, /*dilation_w=*/_dilation,
                                                   CUDNN_CROSS_CORRELATION,
-                                                  data_type));
+                                                  CUDNN_DATA_FLOAT));
+
+      CUDNN_CHECK(cudnnSetConvolutionMathType(conv_desc, CUDNN_DEFAULT_MATH));
+      if (data_type == CUDNN_DATA_HALF)
+        CUDNN_CHECK(cudnnSetConvolutionMathType(conv_desc, CUDNN_TENSOR_OP_MATH));
 
       cudnnHandle_t handle = cuda::get_cudnn_handle();