Can you test it with https://github.com/FluxML/Flux.jl/pull/1402 _Originally posted by @DhairyaLGandhi in https://github.com/aced-differentiate/DeepUncertainty.jl/pull/4#discussion_r700167230_