Merge pull request #7 from tinyzqh/main

fix index of num_features_to_replace bug
shibhansh · Nov 3, 2024 · 63c35f3 · 63c35f3
2 parents 7bf3dfe + d096af7
commit 63c35f3
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/lop/algos/convGnT.py b/lop/algos/convGnT.py
@@ -207,7 +207,7 @@ def update_optim_params(self, features_to_replace_input_indices, features_to_rep
         if self.opt_type == 'AdamGnT':
             for i in range(self.num_hidden_layers):
                 # input weights
-                if num_features_to_replace == 0:
+                if num_features_to_replace[i] == 0:
                     continue
                 # input weights
                 self.opt.state[self.net[i * 2].bias]['exp_avg'][features_to_replace_input_indices[i]] = 0.0

diff --git a/lop/algos/gnt.py b/lop/algos/gnt.py
@@ -211,7 +211,7 @@ def update_optim_params(self, features_to_replace, num_features_to_replace):
         if self.opt_type == 'adam':
             for i in range(self.num_hidden_layers):
                 # input weights
-                if num_features_to_replace == 0:
+                if num_features_to_replace[i] == 0:
                     continue
                 self.opt.state[self.net[i * 2].weight]['exp_avg'][features_to_replace[i], :] = 0.0
                 self.opt.state[self.net[i * 2].bias]['exp_avg'][features_to_replace[i]] = 0.0