chore: add type check

metaopt · XuehaiPan · Sep 11, 2022 · Apr 13, 2022 · Apr 13, 2022 · Apr 13, 2022
commit fa2a38ce68c79e31732ac6a73170dd579081a55a
diff --git a/torchopt/_src/optimizer/base.py b/torchopt/_src/optimizer/base.py
@@ -37,6 +37,9 @@ def __init__(self, params: Iterable[torch.Tensor], impl: GradientTransformation)
                 Note that using ``Optimizer(sgd())`` or ``Optimizer(chain(sgd()))`` is equivalent to
                 :class:`torchopt.SGD`.
         """
+        if not isinstance(impl, GradientTransformation):
+            raise TypeError(f'{impl} (type: {type(impl).__name__}) is not a GradientTransformation')
+
         self.impl = impl
         self.param_groups = []  # type: ignore
         self.param_tree_groups = []  # type: ignore

diff --git a/torchopt/_src/optimizer/func/base.py b/torchopt/_src/optimizer/func/base.py
@@ -48,6 +48,9 @@ def __init__(self, impl: GradientTransformation, *, inplace: bool = False) -> No
             inplace (optional): (default: :data:`False`)
                 The default value of ``inplace`` for each optimization update.
         """
+        if not isinstance(impl, GradientTransformation):
+            raise TypeError(f'{impl} (type: {type(impl).__name__}) is not a GradientTransformation')
+
         self.impl = impl
         self.optim_state = self.__NOT_INITIALIZED
         self.inplace = bool(inplace)

diff --git a/torchopt/_src/optimizer/meta/base.py b/torchopt/_src/optimizer/meta/base.py
@@ -37,6 +37,9 @@ def __init__(self, net: nn.Module, impl: GradientTransformation):
                 ``MetaOptimizer(chain(sgd(moment_requires_grad=True)))`` is equivalent to
                 :class:`torchopt.MetaSGD`.
         """
+        if not isinstance(impl, GradientTransformation):
+            raise TypeError(f'{impl} (type: {type(impl).__name__}) is not a GradientTransformation')
+
         self.impl = impl
         self.param_containers_groups = []  # type: ignore
         self.state_groups = []  # type: ignore