Can't use torch.hann_window in ScriptModule

jit
1

I’m learning jit.script. I got error when I use torch.hann_window.
Code:

import torch
class MyModule(torch.nn.Module):
    def __init__(self):
        super(MyModule, self).__init__()

    def forward(self) -> torch.Tensor:
        return torch.hann_window(5)

if __name__ == "__main__":
    script_module = torch.jit.script(MyModule())
    print(script_module())

The error message is:

Traceback (most recent call last):
  File "D:\Coding\Work\test_jit.py", line 15, in <module>
    print(script_module())
  File "D:\PublicVENV\lib\site-packages\torch\nn\modules\module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
RuntimeError: The following operation failed in the TorchScript interpreter.
Traceback of TorchScript (most recent call last):
  File "D:\Coding\Work\test_jit.py", line 10, in forward
    def forward(self) -> torch.Tensor:
        return torch.hann_window(5, True)
               ~~~~~~~~~~~~~~~~~ <--- HERE
RuntimeError: result type Float can't be cast to the desired output type Long

How can I use hann_window in jit script ? Thanks !!

2

Could you create a GitHub issue for this error, as I would expect torch.hann_window to work in the same way as e.g. torch.randn (the latter works in your code snippet).

3

If we set the dtype at the calling, it could scripted correctly

import torch
class MyModule(torch.nn.Module):
    def __init__(self):
        super(MyModule, self).__init__()

    def forward(self) -> torch.Tensor:
        return torch.hann_window(5,dtype=torch.float32)

if __name__ == "__main__":
    script_module = torch.jit.script(MyModule())
    print(script_module())