able to wrap a bunch of subnetworks and then control whether filtering is on /off, in the case torch func breaks for some large complicated network

Author: lucidrains

GAF_microbatch_pytorch/GAF.py

@@ -224,3 +224,20 @@ def forward(
 
         out = gaf_function(tree_spec, *tree_nodes)
         return out
+
+# helper functions for disabling GAF wrappers within a network
+# for handy ablation, in the case subnetworks within a neural network were wrapped
+
+def set_filter_gradients_(
+    m: Module,
+    filter_gradients: bool,
+    filter_distance_thres = None
+):
+    for module in m.modules():
+        if not isinstance(module, GAFWrapper):
+            continue
+
+        module.filter_gradients = filter_gradients
+
+        if exists(filter_distance_thres):
+            module.filter_distance_thres = filter_distance_thres

GAF_microbatch_pytorch/__init__.py

@@ -1 +1,4 @@
-from GAF_microbatch_pytorch.GAF import GAFWrapper
+from GAF_microbatch_pytorch.GAF import (
+    GAFWrapper,
+    set_filter_gradients_
+)

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "GAF-microbatch-pytorch"
-version = "0.0.2"
+version = "0.0.3"
 description = "Gradient Agreement Filtering"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }