update: remove torch.jit.script decorator

kozistr · Feb 1, 2025 · 68a02f8 · 68a02f8
1 parent 6772208
commit 68a02f8
Showing 1 changed file with 0 additions and 1 deletion.
diff --git a/pytorch_optimizer/optimizer/psgd_utils.py b/pytorch_optimizer/optimizer/psgd_utils.py
@@ -72,7 +72,6 @@ def triu_with_diagonal_and_above(a: torch.Tensor) -> torch.Tensor:
     return torch.triu(a, diagonal=0) + torch.triu(a, diagonal=1)
 
 
-@torch.jit.script
 def update_precondition_dense(
     q: torch.Tensor, dxs: List[torch.Tensor], dgs: List[torch.Tensor], step: float = 0.01, eps: float = 1.2e-38
 ) -> torch.Tensor: