Update cosine

alvgaona · alvgaona · commit 0bb072ff8bf5 · 2022-10-11T14:14:05.000-03:00
diff --git a/torch/signal/windows/__init__.py b/torch/signal/windows/__init__.py
@@ -52,6 +52,10 @@ def add_docstr(function, docstr):
     >>> torch.signal.windows.cosine(10,periodic=False)
     tensor([0.1564, 0.4540, 0.7071, 0.8910, 0.9877, 0.9877, 0.8910, 0.7071, 0.4540,
     0.1564])
+
+.. note::
+    The window is normalized with the maximum value equal to 1, however, the 1 doesn't appear if `M` is even
+    and `periodic` is `False`.
 """.format(
         **factory_common_args
     ),
diff --git a/torch/signal/windows/windows.py b/torch/signal/windows/windows.py
@@ -86,14 +86,14 @@ def cosine(window_length: int,
     if window_length == 1:
         return torch.ones((1,), dtype=dtype, layout=layout, device=device, requires_grad=requires_grad)
 
-    if periodic:
-        window_length += 1
-
     start = .5
-    k = torch.arange(start, start + window_length, dtype=dtype, layout=layout, device=device, requires_grad=requires_grad)
-    window = torch.sin(torch.pi / window_length * k)
-
-    return window[:-1] if periodic else window
+    k = torch.arange(start,
+                     start+window_length,
+                     dtype=dtype,
+                     layout=layout,
+                     device=device,
+                     requires_grad=requires_grad)
+    return torch.sin(torch.pi / (window_length + 1 if periodic else window_length) * k)
 
 
 def gaussian(window_length: int,
