fix(optim): correct interface to _inner

Author: ClashLuke

setup.py

@@ -10,7 +10,7 @@
     name='truegrad',
     license='BSD',
     description='PyTorch interface for TrueGrad-AdamW',
-    version='2.3.4',
+    version='2.3.5',
     long_description=README,
     url='https://github.com/clashluke/truegrad',
     packages=setuptools.find_packages(),

truegrad/optim.py

@@ -92,7 +92,7 @@ def step(self, closure=None):
 
                 step = step_t.item()
 
-                base_update, update, alpha = self._inner(step, p,
+                base_update, update, alpha = self._inner(step, p, group,
                                                          **{k: state.get(k) for k in self.shared_statistics},
                                                          **{k: state.get(k) for k in self.base_statistics},
                                                          **{k: state.get(k) for k in self.true_statistics})
@@ -132,7 +132,7 @@ def __init__(self, params, lr: float = 1e-3,
                          decay_to_init=decay_to_init, default_to_baseline=default_to_baseline,
                          enforce_baseline=enforce_baseline)
 
-    def _inner(self, step: int, p: Parameter, do_baseline: bool, group: Dict[str, Any], exp_avg: Tensor,
+    def _inner(self, step: int, p: Parameter, group: Dict[str, Any], exp_avg: Tensor,
                exp_avg_sq: Optional[Tensor] = None, exp_avg_true_sq: Optional[Tensor] = None
                ) -> Tuple[Optional[Tensor], Optional[Tensor], float]:
         if len(group["betas"]) == 2:
@@ -166,7 +166,7 @@ def __init__(self, params, lr: float = 1e-3,
                          decay_to_init=decay_to_init, default_to_baseline=default_to_baseline,
                          enforce_baseline=enforce_baseline)
 
-    def _inner(self, step: int, p: Parameter, do_baseline: bool, group: Dict[str, Any],
+    def _inner(self, step: int, p: Parameter, group: Dict[str, Any],
                exp_avg: Optional[Tensor] = None, exp_avg_sq: Optional[Tensor] = None,
                exp_avg_true: Optional[Tensor] = None, exp_avg_true_sq: Optional[Tensor] = None
                ) -> Tuple[Optional[Tensor], Optional[Tensor], float]: