Merge pull request #240 from michaldyczko/reg_noise_typing

fix: reg_noise typing
kozistr · May 29, 2024 · 7205a27 · 7205a27
2 parents eda736f + 729a24e
commit 7205a27
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/pytorch_optimizer/optimizer/utils.py b/pytorch_optimizer/optimizer/utils.py
@@ -282,7 +282,7 @@ def reduce_max_except_dim(x: torch.Tensor, dim: int) -> torch.Tensor:
 
 def reg_noise(
     network1: nn.Module, network2: nn.Module, num_data: int, lr: float, eta: float = 8e-3, temperature: float = 1e-4
-) -> torch.Tensor | float:
+) -> Union[torch.Tensor, float]:
     r"""Entropy-MCMC: Sampling from flat basins with ease.
 
     usage: https://github.com/lblaoke/EMCMC/blob/master/exp/cifar10_emcmc.py