fix(optimizer): add closure support (needed for PL)

setup.py

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

truegrad/nn.py

@@ -143,7 +143,7 @@ def forward(ctx, out, fn, args, kwargs) -> torch.Tensor:
     @staticmethod
     def backward(ctx, dy: torch.Tensor) -> Tuple[None, None, None, None]:
         def _square(x: Union[torch.Tensor, None]):
-            if isinstance(x, TrueGradParameter):
+            if isinstance(x, torch.nn.Parameter):
                 x = x.data
             if not isinstance(x, torch.Tensor) or not torch.is_floating_point(x):
                 return x

truegrad/optim.py

@@ -16,6 +16,11 @@ def __init__(self, params, lr: float = 1e-3,
 
     @torch.no_grad()
     def step(self, closure=None):
+        if closure is None:
+            loss = None
+        else:
+            with torch.enable_grad():
+                loss = closure()
         for group in self.param_groups:
             if len(group["betas"]) == 2:
                 beta1, beta2 = group["betas"]
@@ -74,3 +79,4 @@ def step(self, closure=None):
                     alpha = alpha * adam_update.norm() / update.norm()
 
                 p.add_(update, alpha=alpha)
+        return loss