Use CUDA in LayerNorm gradient computation

I don't know how I could miss this.

Use CUDA in LayerNorm gradient computation
I don't know how I could miss this.
49314c12 · Adrià Arrufat · Davis E. King · 3a267db5 · 49314c12
Commit 49314c12 authored Apr 10, 2022 by Adrià Arrufat Committed by Davis E. King Apr 12, 2022
Hide whitespace changes
Inline Side-by-side

Showing with 4 additions and 0 deletions

dlib/cuda/tensor_tools.cpp dlib/cuda/tensor_tools.cpp +4 -0

No files found.
--- a/dlib/cuda/tensor_tools.cpp
+++ b/dlib/cuda/tensor_tools.cpp
@@ -687,7 +687,11 @@ namespace dlib { namespace tt
            tensor& beta_grad
    )
    {
+#ifdef DLIB_USE_CUDA
+        cuda::layer_normalize_gradient(eps, gradient_input, means, invstds, src, gamma, src_grad, gamma_grad, beta_grad);
+#else
        cpu::layer_normalize_gradient(eps, gradient_input, means, invstds, src, gamma, src_grad, gamma_grad, beta_grad);
+#endif
    }

 // ----------------------------------------------------------------------------------------