fix missing import Optional

reference_algorithms/paper_baselines/momentum/jax/submission.py

@@ -1,7 +1,7 @@
 """Submission file for a SGD with HeavyBall momentum optimizer in Jax."""
 
 import functools
-from typing import Any, Callable, Dict, Iterator, List, Tuple
+from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple
 
 from flax import jax_utils
 import jax

reference_algorithms/paper_baselines/momentum/pytorch/submission.py

@@ -1,6 +1,6 @@
 """Submission file for a SGD with HeavyBall momentum optimizer in PyTorch."""
 
-from typing import Any, Callable, Dict, Iterator, List, Tuple
+from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple
 
 from absl import logging
 import optax

reference_algorithms/paper_baselines/nesterov/jax/submission.py

@@ -1,7 +1,7 @@
 """Submission file for a SGD with Nesterov momentum optimizer in Jax."""
 
 import functools
-from typing import Any, Callable, Dict, Iterator, List, Tuple
+from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple
 
 from flax import jax_utils
 import jax

reference_algorithms/paper_baselines/nesterov/pytorch/submission.py

@@ -1,6 +1,6 @@
 """Submission file for a SGD with Nesterov momentum optimizer in PyTorch."""
 
-from typing import Any, Callable, Dict, Iterator, List, Tuple
+from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple
 
 from absl import logging
 import optax

reference_algorithms/paper_baselines/sam/pytorch/submission.py

@@ -1,6 +1,6 @@
 """Submission file for a SAM optimizer with warmup+cosine LR in PyTorch."""
 
-from typing import Any, Callable, Dict, Iterator, List, Tuple
+from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple
 
 from absl import logging
 import torch

reference_algorithms/target_setting_algorithms/jax_submission_base.py

@@ -1,6 +1,6 @@
 """Update submission function in Jax."""
 import functools
-from typing import Any, Dict, List, Tuple
+from typing import Any, Dict, List, Optional, Tuple
 
 import jax
 from jax import lax

reference_algorithms/target_setting_algorithms/pytorch_submission_base.py

@@ -1,6 +1,6 @@
 """Batch size and update submission functions in PyTorch."""
 
-from typing import Any, Dict, List, Tuple
+from typing import Any, Dict, List, Optional, Tuple
 
 from absl import logging
 import torch