Skip to content

Commit 729a24e

Browse files
committed
fix: reg_noise typing
1 parent eda736f commit 729a24e

File tree

1 file changed

+1
-1
lines changed

1 file changed

+1
-1
lines changed

pytorch_optimizer/optimizer/utils.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -282,7 +282,7 @@ def reduce_max_except_dim(x: torch.Tensor, dim: int) -> torch.Tensor:
282282

283283
def reg_noise(
284284
network1: nn.Module, network2: nn.Module, num_data: int, lr: float, eta: float = 8e-3, temperature: float = 1e-4
285-
) -> torch.Tensor | float:
285+
) -> Union[torch.Tensor, float]:
286286
r"""Entropy-MCMC: Sampling from flat basins with ease.
287287
288288
usage: https://github.com/lblaoke/EMCMC/blob/master/exp/cifar10_emcmc.py

0 commit comments

Comments
 (0)