🐛 Bug

To Reproduce

Steps to reproduce the behavior:

1. use torch.chunk on torch.jit.script_method

import torch

class BatchMod2(torch.jit.ScriptModule):
    def __init__(self, x_dim, z_dim):
        super().__init__()

    @torch.jit.script_method
    def forward(self, x, z):
        mean_scale = z.repeat(1, 2)
        mean, scale = torch.chunk(mean_scale[..., None, None], 2, 1)
        y = x * (scale + 1) + mean
        return y


b = BatchMod2(6, 6)
a1 = torch.ones(2,6,4,4)
a2 = torch.ones(2,6)
c=b(a1, a2)

Expected behavior

No exception.

Environment

Please copy and paste the output from our
environment collection script
(or fill out the checklist below manually).

You can get the script and run it with:

wget https://raw.githubusercontent.com/pytorch/pytorch/master/torch/utils/collect_env.py
# For security purposes, please check the contents of collect_env.py before running it.
python collect_env.py

• PyTorch Version (e.g., 1.0): 1.1.0
• OS (e.g., Linux): win10
• How you installed PyTorch (conda, pip, source): pip
• Build command you used (if compiling from source):
• Python version: 3.7.1
• CUDA/cuDNN version: 10.0
• GPU models and configuration:gtx970m
• Any other relevant information:

Additional context

When I comment @torch.jit.script_method, it can work normally.