# Modifying Flux source code for GPU
**URL:** https://discourse.julialang.org/t/modifying-flux-source-code-for-gpu/96360
**Category:** Specific Domains
**Created:** [March 20, 2023, 7:16pm UTC](https://discourse.julialang.org/t/modifying-flux-source-code-for-gpu/96360 "2023-03-20T19:16:43Z")
**Posts on this page:** 5
**Page:** 1
### Author:  [@wsshin](https://discourse.julialang.org/u/wsshin)
#### Post date: [March 20, 2023, 7:16pm UTC](https://discourse.julialang.org/t/modifying-flux-source-code-for-gpu/96360/1 "2023-03-20T19:16:43Z")
In Flux.jl, I wanted to modify the behavior of `BatchNorm`, so I changed the code [here](https://github.com/FluxML/Flux.jl/blob/c850df5409ca545be433dec835034cffa8486aa4/src/layers/normalise.jl). The modified code was effective when it ran on CPU, but I realized that it was not effective on GPU.
After some investigation, I found out that that after transferring the model `m` to GPU by `m = m |> gpu`, the above linked code was no longer executed.
What is the general procedure to make the changes in the Flux.jl source code effective on GPU as well? I think the actual code executed on GPU might be [this](https://github.com/JuliaGPU/CUDA.jl/blob/b60f37a0b1586be8920a33f0c17d924f8ec58336/lib/cudnn/src/libcudnn.jl#L1350-L1357), but I’m not sure how to modify it because it eventually uses `@ccall`.
---
### Author:  [@ToucheSir](https://discourse.julialang.org/u/ToucheSir)
#### Post date: [March 20, 2023, 9:21pm UTC](https://discourse.julialang.org/t/modifying-flux-source-code-for-gpu/96360/2 "2023-03-20T21:21:34Z")
For historical reasons, the CUDA path for Flux’s batchnorm lives in its [own file](https://github.com/FluxML/Flux.jl/blob/master/src/cuda/cudnn.jl). Thus you’d either have to remove this code or change those methods to make your own version work on GPU. Medium-term, we plan to [clean up](https://github.com/FluxML/NNlib.jl/pull/452) the API layering around norm layers so that all batchnorm layer methods can live in one place.
---
### Author:  [@wsshin](https://discourse.julialang.org/u/wsshin)
#### Post date: [March 25, 2023, 1:51am UTC](https://discourse.julialang.org/t/modifying-flux-source-code-for-gpu/96360/3 "2023-03-25T01:51:00Z")
Thanks. So I was able to track down that the following happens during training:
- `(BN::BatchNorm)(x::CuArray)` calls [`NNlibCUDA.batchnorm()`](https://github.com/FluxML/Flux.jl/blob/master/src/cuda/cudnn.jl#L9-L11)
- `NNlibCUDA.batchnorm()` calls [`cudnnBNForward!()`](https://github.com/FluxML/NNlibCUDA.jl/blob/master/src/cudnn/batchnorm.jl#L37)
- `cudnnBNForward!()` calls [cuDNN.cudnnBatchNormalizationForwardTraining()](https://github.com/FluxML/NNlibCUDA.jl/blob/master/src/cudnn/batchnorm.jl#L80)
- `cuDNN.cudnnBatchNormalizationForwardTraining()` calls [`@ccall libcudnn.cudnnBatchNormalizationForwardTraining()`](https://github.com/JuliaGPU/CUDA.jl/blob/b60f37a0b1586be8920a33f0c17d924f8ec58336/lib/cudnn/src/libcudnn.jl#L1359)
Now, where can I see the code for the final function, `libcudnn.cudnnBatchNormalizationForwardTraining()`?
---
### Author:  [@wsshin](https://discourse.julialang.org/u/wsshin)
#### Post date: [March 27, 2023, 8:51pm UTC](https://discourse.julialang.org/t/modifying-flux-source-code-for-gpu/96360/4 "2023-03-27T20:51:07Z")
Maybe @maleadt could help?
---
### Author:  [@ToucheSir](https://discourse.julialang.org/u/ToucheSir)
#### Post date: [March 27, 2023, 8:57pm UTC](https://discourse.julialang.org/t/modifying-flux-source-code-for-gpu/96360/5 "2023-03-27T20:57:54Z")
The final function is in cuDNN itself, which is closed source. Your best bet for understanding what it does is reading the relevant docs at [NVIDIA Deep Learning cuDNN Documentation](https://docs.nvidia.com/deeplearning/cudnn/).