diff --git a/prior_wd_optim.py b/prior_wd_optim.py
index a827e6e..ef46aee 100644
--- a/prior_wd_optim.py
+++ b/prior_wd_optim.py
@@ -26,7 +26,7 @@ def step(self, closure=None):
         if self.use_prior_wd:
             for i, group in enumerate(self.param_groups):
                 for p in group["params"]:
-                    if self.exclude_last_group and i == len(self.param_groups):
+                    if self.exclude_last_group and i == len(self.param_groups)-1:
                         p.data.add_(-group["lr"] * self.weight_decay_by_group[i], p.data)
                     else:
                         p.data.add_(