fix: [pre-commit.ci] auto fixes [...]

Author: pre-commit-ci[bot]

torchopt/distributed/api.py

@@ -318,12 +318,12 @@ def remote_async_call(
         futures.append(fut)
 
     future = cast(
-        Future[List[T]],
+        'Future[List[T]]',
         torch.futures.collect_all(futures).then(lambda fut: [f.wait() for f in fut.wait()]),
     )
     if reducer is not None:
         return cast(
-            Future[U],
+            'Future[U]',
             future.then(lambda fut: reducer(fut.wait())),
         )
     return future

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)

torchopt/utils.py

@@ -79,13 +79,13 @@ def fn_(obj: Any) -> None:
             obj.detach_().requires_grad_(requires_grad)
 
     if isinstance(target, ModuleState):
-        true_target = cast(TensorTree, (target.params, target.buffers))
+        true_target = cast('TensorTree', (target.params, target.buffers))
     elif isinstance(target, nn.Module):
-        true_target = cast(TensorTree, tuple(target.parameters()))
+        true_target = cast('TensorTree', tuple(target.parameters()))
     elif isinstance(target, MetaOptimizer):
-        true_target = cast(TensorTree, target.state_dict())
+        true_target = cast('TensorTree', target.state_dict())
     else:
-        true_target = cast(TensorTree, target)  # tree of tensors
+        true_target = cast('TensorTree', target)  # tree of tensors
 
     pytree.tree_map_(fn_, true_target)
 
@@ -325,7 +325,7 @@ def recover_state_dict(
     from torchopt.optim.meta.base import MetaOptimizer
 
     if isinstance(target, nn.Module):
-        params, buffers, *_ = state = cast(ModuleState, state)
+        params, buffers, *_ = state = cast('ModuleState', state)
         params_containers, buffers_containers = extract_module_containers(target, with_buffers=True)
 
         if state.detach_buffers:
@@ -343,7 +343,7 @@ def clone_detach_(t: torch.Tensor) -> torch.Tensor:
         ):
             tgt.update(src)
     elif isinstance(target, MetaOptimizer):
-        state = cast(Sequence[OptState], state)
+        state = cast('Sequence[OptState]', state)
         target.load_state_dict(state)
     else:
         raise TypeError(f'Unexpected class of {target}')
@@ -422,9 +422,9 @@ def module_clone(  # noqa: C901
 
     if isinstance(target, (nn.Module, MetaOptimizer)):
         if isinstance(target, nn.Module):
-            containers = cast(TensorTree, extract_module_containers(target, with_buffers=True))
+            containers = cast('TensorTree', extract_module_containers(target, with_buffers=True))
         else:
-            containers = cast(TensorTree, target.state_dict())
+            containers = cast('TensorTree', target.state_dict())
         tensors = pytree.tree_leaves(containers)
         memo = {id(t): t for t in tensors}
         cloned = copy.deepcopy(target, memo=memo)
@@ -476,7 +476,7 @@ def clone_detach_(t: torch.Tensor) -> torch.Tensor:
     else:
         replicate = clone_detach_
 
-    return pytree.tree_map(replicate, cast(TensorTree, target))
+    return pytree.tree_map(replicate, cast('TensorTree', target))
 
 
 @overload

torchopt/visual.py

@@ -129,7 +129,7 @@ def make_dot(  # noqa: C901
                 elif isinstance(param, Generator):
                     param_map.update({v: k for k, v in param})
                 else:
-                    param_map.update({v: k for k, v in cast(Mapping, param).items()})
+                    param_map.update({v: k for k, v in cast('Mapping', param).items()})
 
     node_attr = {
         'style': 'filled',