Update pf.py

hansenlab · Oct 9, 2024 · f5cb86b · f5cb86b
1 parent e167e95
commit f5cb86b
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/mNSF/NSF/pf.py b/mNSF/NSF/pf.py
@@ -25,6 +25,7 @@
 dtp = "float32"
 rng = np.random.default_rng()
 
+@tf.custom_gradient
 def checkpoint_grad(x):
     y = tf.identity(x)
     def grad(dy):