Repository: pytorch
Issue: Pytorch mobile Vulkan API bug during call to ReLU
Labels: oncall: mobile,oncall: java

## 🐛 Bug

I have a simple model in pytorch-nightly build (1.8.0a0+bdf360f). I traced it and try to run on android, using Java API.
It works on default device (CPU), but crashes on Vulkan (GPU). AARs for android are generated from the following commit f5b68e74d75c38d5e2044fc6b62112181080bb3f.

It reports: Tried to call KernelFunction::callBoxed() on a KernelFunction that can only be called with KernelFunction::call(). opname: aten::relu


<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->
Pytorch code to create model:
```
class reluwrapper(nn.Module):
    def __init__(self):
        super().__init__()
        self.r1=torch.nn.ReLU()
        
    def forward(self, x):
        outputs=self.r1(x)
        
        return outputs
```

Code to trace model:

```
model=reluwrapper()
inputs=torch.randn(1,192,4,4)
export=True
if export:
    with torch.no_grad():        
        res = model(inputs)
        traced_script_module = torch.jit.trace(model, (inputs))
        traced_script_module.save("assets/relu.pt")
```


As a result, if I call it within java android app, I get the following trace:

```
 E/AndroidRuntime: FATAL EXCEPTION: pool-4-thread-1
    Process: *****, PID: 19220
    java.lang.RuntimeException: The following operation failed in the TorchScript interpreter.
    Traceback of TorchScript, serialized code (most recent call last):
      File "code/__torch__.py", line 8, in forward
      def forward(self: __torch__.reluwrapper,
        input: Tensor) -> Tensor:
        return (self.r1).forward(input, )
                ~~~~~~~~~~~~~~~~ <--- HERE
      File "code/__torch__/torch/nn/modules/activation/___torch_mangle_1168.py", line 7, in forward
      def forward(self: __torch__.torch.nn.modules.activation.___torch_mangle_1168.ReLU,
        input: Tensor) -> Tensor:
        return torch.relu(input)
               ~~~~~~~~~~ <--- HERE
    
    Traceback of TorchScript, original code (most recent call last):
    /home/alex/.conda/envs/torch-nightly/lib/python3.7/site-packages/torch/nn/functional.py(1111): relu
    /home/alex/.conda/envs/torch-nightly/lib/python3.7/site-packages/torch/nn/modules/activation.py(102): forward
    /home/alex/.conda/envs/torch-nightly/lib/python3.7/site-packages/torch/nn/modules/module.py(726): _slow_forward
    /home/alex/.conda/envs/torch-nightly/lib/python3.7/site-packages/torch/nn/modules/module.py(742): _call_impl
    <ipython-input-30-034594f31033>(7): forward
    /home/alex/.conda/envs/torch-nightly/lib/python3.7/site-packages/torch/nn/modules/module.py(726): _slow_forward
    /home/alex/.conda/envs/torch-nightly/lib/python3.7/site-packages/torch/nn/modules/module.py(742): _call_impl
    /home/alex/.conda/envs/torch-nightly/lib/python3.7/site-packages/torch/jit/_trace.py(940): trace_module
    /home/alex/.conda/envs/torch-nightly/lib/python3.7/site-packages/torch/jit/_trace.py(742): trace
    <ipython-input-34-99ba286aad47>(7): <module>
    /home/alex/.conda/envs/torch-nightly/lib/python3.7/site-packages/IPython/core/interactiveshell.py(3418): run_code
    /home/alex/.conda/envs/torch-nightly/lib/python3.7/site-packages/IPython/core/interactiveshell.py(3338): run_ast_nodes
    /home/alex/.conda/envs/torch-nightly/lib/python3.7/site-packages/IPython/core/interactiveshell.py(3147): run_cell_async
    /home/alex/.conda/envs/torch-nightly/lib/python3.7/site-packages/IPython/core/async_helpers.py(68): _pseudo_sync_runner
    /home/alex/.conda/envs/torch-nightly/lib/python3.7/site-packages/IPython/core/interactiveshell.py(2923): _run_cell
    /home/alex/.conda/envs/torch-nightly/lib/python3.7/site-packages/IPython/core/interactiveshell.py(2878): run_cell
    /home/alex/.conda/envs/torch-nightly/lib/python3.7/site-packages/ipykernel/zmqshell.py(536): run_cell
    /home/alex/.conda/envs/torch-nightly/lib/python3.7/site-packages/ipykernel/ipkernel.py(306): do_execute
    /home/alex/.conda/envs/torch-nightly/lib/python3.7/site-packages/tornado/gen.py(234): wrapper
    /home/alex/.conda/envs/torch-nightly/lib/python3.7/site-packages/ipykernel/kernelbase.py(545): execute_request
    /home/alex/.conda/envs/torch-nightly/lib/python3.7/site-packages/tornado/gen.py(234): wrapper
    /home/alex/.conda/envs/torch-nightly/lib/python3.7/site-packages/ipykernel/kernelbase.py(268): dispatch_shell
    /home/alex/.conda/envs/torch-nightly/lib/python3.7/site-packages/tornado/gen.py(234): wrapper
    /home/alex/.conda/envs/torch-nightly/lib/python3.7/site-packages/ipykernel/kernelbase.py(365): process_one
    /home/alex/.conda/envs/torch-nightly/lib/python3.7/site-packages/tornado/gen.py(775): run
    /home/alex/.conda/envs/torch-nightly/lib/python3.7/site-packages/tornado/gen.py(814): inner
    /home/alex/.conda/envs/torch-nightly/lib/python3.7/site-packages/tornado/ioloop.py(741): _run_callback
    /home/alex/.conda/envs/torch-nightly/lib/python3.7/site-packages/tornado/ioloop.py(688): <lambda>
    /home/alex/.conda/envs/torch-nightly/lib/python3.7/asyncio/events.py(88): _run
    /home/alex/.conda/envs/torch-nightly/lib/python3.7/asyncio/base_events.py(1786): _run_once
    /home/alex/.conda/envs/torch-nightly/lib/python3.7/asyncio/base_events.py(541): run_forever
    /home/alex/.conda/envs/torch-nightly/lib/python3.7/site-packages/tornado/platform/asyncio.py(199): start
E/AndroidRuntime: /home/alex/.conda/envs/torch-nightly/lib/python3.7/site-packages/ipykernel/kernelapp.py(612): start
    /home/alex/.conda/envs/torch-nightly/lib/python3.7/site-packages/traitlets/config/application.py(845): launch_instance
    /home/alex/.conda/envs/torch-nightly/lib/python3.7/site-packages/ipykernel_launcher.py(16): <module>
    /home/alex/.conda/envs/torch-nightly/lib/python3.7/runpy.py(85): _run_code
    /home/alex/.conda/envs/torch-nightly/lib/python3.7/runpy.py(193): _run_module_as_main
    RuntimeError: false INTERNAL ASSERT FAILED at "../aten/src/ATen/core/boxing/KernelFunction.cpp":76, please report a bug to PyTorch. Tried to call KernelFunction::callBoxed() on a KernelFunction that can only be called with KernelFunction::call(). opname: aten::relu If you're using mobile selective build please make sure to include all ops exported from `torch.jit.export_opnames(model)`.
    
        at org.pytorch.NativePeer.forward(Native Method)
        at org.pytorch.Module.forward(Module.java:49)
        at com.example.****.MainActivity$NNThread.run(MainActivity.java:584)
        at java.util.concurrent.ThreadPoolExecutor.runWorker(ThreadPoolExecutor.java:1167)
        at java.util.concurrent.ThreadPoolExecutor$Worker.run(ThreadPoolExecutor.java:641)
        at java.lang.Thread.run(Thread.java:919)

```
## Expected behavior

It is expected to run the module and not to crash

## Environment

 - PyTorch Version (e.g., 1.0): 1.8.0a0+bdf360f for export,  f5b68e74d75c38d5e2044fc6b62112181080bb3f for pytorch mobile
 - OS (e.g., Linux): Ubuntu 18.04
 - How you installed PyTorch (`conda`, `pip`, source): source
 - Build command you used (if compiling from source): 
export CMAKE_PREFIX_PATH=${CONDA_PREFIX:-"$(dirname $(which conda))/../"}
python setup.py install
 - Python version: 3.7
 - CUDA/cuDNN version: 11.1
 - GPU models and configuration:
 - Any other relevant information: