The type issue of jit.script

Hello, I have some question about the type issue of torch.jit.script. As I run the example code, it always reports on line 104 that the third input variable should be Tensor. But I already set the type at line 68-line 69

How should I solve this problem or is there any work-around?

have you solved it?
I met the similar problem.
The default variable in jit script is Tensor, but I need int or .cpu().numpy().
If you have any hints, please inform me, thx a lot!