OrgLion-ML
/
dlib
mirror of https://github.com/davisking/dlib.git
1
0
Fork 0
Code Issues Releases Wiki Activity

Use CUDA in LayerNorm gradient computation 

I don't know how I could miss this.
This commit is contained in:
Adrià Arrufat 2022-04-10 22:31:22 +09:00 committed by Davis E. King
parent 3a267db577
commit 49314c12d9
1 changed files with 4 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
dlib/cuda/tensor_tools.cpp
View File

@ -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
}
// ----------------------------------------------------------------------------------------