fix: optimiser typing

MarkTuddenham · Oct 1, 2021 · 68a14b4 · 68a14b4
1 parent 219456b
commit 68a14b4
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/orth_optim/__init__.py b/src/orth_optim/__init__.py
@@ -6,7 +6,7 @@
 
 import torch
 
-OptimType = TypeVar('U', bound=torch.optim.Optimizer)
+OptimType = TypeVar('OptimType', bound=torch.optim.Optimizer)
 
 logger: logging.Logger = logging.getLogger(__name__)
 logger.setLevel(logging.DEBUG)
@@ -31,7 +31,7 @@ def _orth_grads(optimiser: OptimType) -> None:
                             G,
                             q=1,    # assume rank is at least 1
                             M=1e-4 * G.mean() * torch.randn_like(G))
-                        orth_G: torch.Tensor = u @ v.T
+                        orth_G = u @ v.T
                     except RuntimeError:
                         logger.error(('Failed to perform SVD with noise,'
                                       ' skipping gradient orthogonalisation'))