From f85efec683912d6f22c5a4149af1cdf057f99cf1 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Mon, 5 Aug 2024 20:15:56 +0000
Subject: [PATCH] fix: [pre-commit.ci] auto fixes [...]

---
 torchopt/nn/stateless.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/torchopt/nn/stateless.py b/torchopt/nn/stateless.py
index c7f92b86..0f4f17b3 100644
--- a/torchopt/nn/stateless.py
+++ b/torchopt/nn/stateless.py
@@ -84,7 +84,7 @@ def reparametrize(
     module: nn.Module,
     named_tensors: dict[str, torch.Tensor] | Iterable[tuple[str, torch.Tensor]],
     allow_missing: bool = False,
-) -> Generator[nn.Module, None, None]:
+) -> Generator[nn.Module]:
     """Reparameterize the module parameters and/or buffers."""
     if not isinstance(named_tensors, dict):
         named_tensors = dict(named_tensors)