fixed other mypy bugs

cmu-db · Sep 3, 2024 · 10cbfe1 · 10cbfe1
1 parent 31d4330
commit 10cbfe1
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/tune/protox/agent/build_trial.py b/tune/protox/agent/build_trial.py
@@ -126,8 +126,9 @@ def _modify_benchbase_config(
 
 def _gen_noise_scale(
  vae_config: dict[str, Any], hpo_params: dict[str, Any]
-) -> Callable[[ProtoAction, torch.Tensor], ProtoAction]:
- def f(p: ProtoAction, n: torch.Tensor) -> ProtoAction:
+) -> Callable[[ProtoAction, Optional[torch.Tensor]], ProtoAction]:
+ def f(p: ProtoAction, n: Optional[torch.Tensor]) -> ProtoAction:
+ assert n is not None
  if hpo_params["scale_noise_perturb"]:
  return ProtoAction(
  torch.clamp(

diff --git a/tune/protox/env/space/latent_space/lsc_index_space.py b/tune/protox/env/space/latent_space/lsc_index_space.py
@@ -35,7 +35,7 @@ def __init__(
  latent_dim: int = 0,
  index_output_transform: Optional[Callable[[ProtoAction], ProtoAction]] = None,
  index_noise_scale: Optional[
- Callable[[ProtoAction, torch.Tensor], ProtoAction]
+ Callable[[ProtoAction, Optional[torch.Tensor]], ProtoAction]
  ] = None,
  logger: Optional[Logger] = None,
  lsc: Optional[LSC] = None,