Repository: pytorch
Issue: [PyTorch Mobile] Can't use Vulkan backend
Labels: triaged,module: android,oncall: mobile,module: vulkan

## 🐛 Bug

<!-- A clear and concise description of what the bug is. -->

## To Reproduce

Steps to reproduce the behavior:

Trying out the vulkan [prototype](https://github.com/pytorch/tutorials/blob/master/prototype_source/vulkan_workflow.rst). The guide is outdated (e.g., installMbvulkanLocalBaseDebug task doesn't even exist)
1. compile aar using:
`USE_VULKAN=1 sh ./scripts/build_pytorch_android.sh arm64-v8a`
2. import 2 aar's and use 'mModule = Module.load(moduleFileAbsoluteFilePath, Device.VULKAN);' in [PyTorchDemoApp](https://github.com/pytorch/android-demo-app/tree/master/PyTorchDemoApp)
3. generate model using:
```
model = torchvision.models.mobilenet_v2(pretrained=True)
model.eval()
script_model = torch.jit.script(model)
from torch.utils.mobile_optimizer import optimize_for_mobile
script_model_vulkan = optimize_for_mobile(script_model, backend='Vulkan')
torch.jit.save(script_model_vulkan, "mobilenet2-vulkan.pt")
```
And replace 'resnet18.pt' w/ the vulkan model
4. run the app

<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->
Here is the stacktrace:
```
2021-01-19 20:48:32.326 13323-13496/org.pytorch.demo E/PyTorchDemo: Error during image analysis
    java.lang.RuntimeException: The following operation failed in the TorchScript interpreter.
    Traceback of TorchScript, serialized code (most recent call last):
      File "code/__torch__/___torch_mangle_483.py", line 30, in forward
        _0 = torch.contiguous(x, memory_format=2)
        _1 = getattr(self, "prepack_folding._jit_pass_packed_weight_0")
        _2 = ops.prepacked.conv2d_clamp_run(_0, _1)
             ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
        x0 = torch.max_pool2d(_2, [3, 3], [2, 2], [1, 1], [1, 1], False)
        _3 = getattr(self, "prepack_folding._jit_pass_packed_weight_1")
    
    Traceback of TorchScript, original code (most recent call last):
                %weight, %bias, %stride, %padding, %dilation, %groups,
                %output_min, %output_max)
            %r = prepacked::conv2d_clamp_run(%input, %packed_weight_bias)
                 ~~~~~~~~~ <--- HERE
            return (%r) 
    RuntimeError: Could not run 'prepacked::conv2d_clamp_run' with arguments from the 'Vulkan' backend. This could be because the operator doesn't exist for this backend, or was omitted during the selective/custom build process (if using custom build). If you are a Facebook employee using PyTorch on mobile, please visit https://fburl.com/ptmfixes for possible resolutions. 'prepacked::conv2d_clamp_run' is only available for these backends: [CPU, BackendSelect, Named, AutogradOther, AutogradCPU, AutogradCUDA, AutogradXLA, Tracer, Autocast, Batched, VmapMode].
    
    CPU: registered at /usr/local/home/wind/Desktop/pytorch-src/aten/src/ATen/native/xnnpack/RegisterOpContextClass.cpp:84 [kernel]
    BackendSelect: fallthrough registered at /usr/local/home/wind/Desktop/pytorch-src/aten/src/ATen/core/BackendSelectFallbackKernel.cpp:3 [backend fallback]
    Named: registered at /usr/local/home/wind/Desktop/pytorch-src/aten/src/ATen/core/NamedRegistrations.cpp:7 [backend fallback]
    AutogradOther: fallthrough registered at /usr/local/home/wind/Desktop/pytorch-src/aten/src/ATen/core/VariableFallbackKernel.cpp:35 [backend fallback]
    AutogradCPU: fallthrough registered at /usr/local/home/wind/Desktop/pytorch-src/aten/src/ATen/core/VariableFallbackKernel.cpp:39 [backend fallback]
    AutogradCUDA: fallthrough registered at /usr/local/home/wind/Desktop/pytorch-src/aten/src/ATen/core/VariableFallbackKernel.cpp:43 [backend fallback]
    AutogradXLA: fallthrough registered at /usr/local/home/wind/Desktop/pytorch-src/aten/src/ATen/core/VariableFallbackKernel.cpp:47 [backend fallback]
    Tracer: fallthrough registered at /usr/local/home/wind/Desktop/pytorch-src/torch/csrc/jit/frontend/tracer.cpp:995 [backend fallback]
    Autocast: fallthrough registered at /usr/local/home/wind/Desktop/pytorch-src/aten/src/ATen/autocast_mode.cpp:250 [backend fallback]
    Batched: registered at /usr/local/home/wind/Desktop/pytorch-src/aten/src/ATen/BatchingRegistrations.cpp:1011 [backend fallback]
    VmapMode: fallthrough registered at /usr/local/home/wind/Desktop/pytorch-src/aten/src/ATen/VmapModeRegistrations.cpp:33 [backend fallback]
    
    
        at org.pytorch.NativePeer.forward(Native Method)
        at org.pytorch.Module.forward(Module.java:49)
        at org.pytorch.demo.vision.ImageClassificationActivity.analyzeImage(ImageClassificationActivity.java:183)
        at org.pytorch.demo.vision.ImageClassificationActivity.analyzeImage(ImageClassificationActivity.java:32)
        at org.pytorch.demo.vision.AbstractCameraXActivity.lambda$setupCameraX$2$AbstractCameraXActivity(AbstractCameraXActivity.java:90)
        at org.pytorch.demo.vision.-$$Lambda$AbstractCameraXActivity$t0OjLr-l_M0-_0_dUqVE4yqEYnE.analyze(Unknown Source:2)
        at androidx.camera.core.ImageAnalysisAbstractAnalyzer.analyzeImage(ImageAnalysisAbstractAnalyzer.java:57)
        at androidx.camera.core.ImageAnalysisNonBlockingAnalyzer$1.run(ImageAnalysisNonBlockingAnalyzer.java:135)
        at android.os.Handler.handleCallback(Handler.java:883)
        at android.os.Handler.dispatchMessage(Handler.java:100)
        at android.os.Looper.loop(Looper.java:214)
2021-01-19 20:48:32.326 13323-13496/org.pytorch.demo E/PyTorchDemo:     at android.os.HandlerThread.run(HandlerThread.java:67)
```

## Expected behavior
Shouldn't crash
<!-- A clear and concise description of what you expected to happen. -->

## Environment

Please copy and paste the output from our
[environment collection script](https://raw.githubusercontent.com/pytorch/pytorch/master/torch/utils/collect_env.py)
(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): c458558334e4aef5107bb60807c76ea9a7078727
 - OS (e.g., Linux): Android
 - How you installed PyTorch (`conda`, `pip`, source):
 - Build command you used (if compiling from source):
 - Python version:
 - CUDA/cuDNN version:
 - GPU models and configuration:
 - Any other relevant information:

## Additional context

<!-- Add any other context about the problem here. -->