make torch_ppc a torch jit script function

RichieHakim · May 7, 2024 · 1ffcc22 · 1ffcc22
1 parent f0671e7
commit 1ffcc22
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/bnpm/spectral.py b/bnpm/spectral.py
@@ -906,6 +906,7 @@ def ppc(phases, axis=None):
     return ((cosSum**2 + sinSum**2) - N) / (N * (N - 1))
 
 
+@torch.jit.script
 def torch_ppc(phases: torch.Tensor, axis: Optional[List[int]] = None):
     """
     Exactly the same as ``ppc`` but works with torch.jit.script.