KAG is modified

cometscome · Jun 3, 2024 · 4c09ee3 · 4c09ee3
1 parent 5a100ae
commit 4c09ee3
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/KAGnet.jl b/src/KAGnet.jl
@@ -68,7 +68,7 @@ function ChainRulesCore.rrule(::typeof(rdf_foward), x, num_grids, grids, denomin
         for n = 1:length(ybar)
             dLdGdx .+= dy[n] .* ybar[n]
         end
-        return sbar, dLdGdx, sbar, sbar, sbar
+        return sbar, dLdGdx, sbar, -dLdGdx, sbar
     end
     return y, pullback
 end