Repository: pytorch
Issue: [PyTorch Mobile] Android speed benchmark binary crashes when using vulkan
Labels: oncall: mobile

## 🐛 Bug

<!-- A clear and concise description of what the bug is. -->
Not sure this is expected given Vulkan support is still prototype, but when trying out the benchmark binary using Vulkan, it crashes. Running on CPU is fine. The model is converted by:

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")

## To Reproduce

`(experimental) wind-macbookpro2:Vulkan wind$ adb shell "/data/local/tmp/speed_benchmark_torch --model=/data/local/tmp/model.pt" --input_dims="1,3,224,224" --input_type="float" --vulkan=true
Starting benchmark.
Running warmup runs.
Main runs.
terminating with uncaught exception of type std::runtime_error: The following operation failed in the TorchScript interpreter.
Traceback of TorchScript, serialized code (most recent call last):
  File "code/__torch__/torchvision/models/mobilenet.py", line 10, in forward
  def forward(self: __torch__.torchvision.models.mobilenet.MobileNetV2,
    x: Tensor) -> Tensor:
    return (self)._forward_impl(x, )
            ~~~~~~~~~~~~~~~~~~~ <--- HERE
  def _forward_impl(self: __torch__.torchvision.models.mobilenet.MobileNetV2,
    x: Tensor) -> Tensor:
  File "code/__torch__/torchvision/models/mobilenet.py", line 14, in _forward_impl
    x: Tensor) -> Tensor:
    _0 = __torch__.torch.nn.functional.adaptive_avg_pool2d
    x0 = (self.features).forward(x, )
          ~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
    x1 = torch.reshape(_0(x0, [1, 1], ), [(torch.size(x0))[0], -1])
    return (self.classifier).forward(x1, )
  File "code/__torch__/torch/nn/modules/container/___torch_mangle_80.py", line 45, in forward
    _17 = getattr(self, "17")
    _18 = getattr(self, "18")
    input0 = (_0).forward(input, )
              ~~~~~~~~~~~ <--- HERE
    input1 = (_1).forward(input0, )
    input2 = (_2).forward(input1, )
  File "code/__torch__/torchvision/models/mobilenet.py", line 29, in forward
    _2 = getattr(self, "1")
    _3 = getattr(self, "2")
    input0 = (_1).forward(input, )
              ~~~~~~~~~~~ <--- HERE
    input1 = (_2).forward(input0, )
    return (_3).forward(input1, )
  File "code/__torch__/torch/nn/modules/conv.py", line 20, in forward
  def forward(self: __torch__.torch.nn.modules.conv.Conv2d,
    input: Tensor) -> Tensor:
    _0 = (self)._conv_forward(input, self.weight, )
          ~~~~~~~~~~~~~~~~~~~ <--- HERE
    return _0
  def _conv_forward(self: __torch__.torch.nn.modules.conv.Conv2d,
  File "code/__torch__/torch/nn/modules/conv.py", line 25, in _conv_forward
    input: Tensor,
    weight: Tensor) -> Tensor:
    _1 = torch.conv2d(input, weight, self.bias, [2, 2], [1, 1], [1, 1], 1)
         ~~~~~~~~~~~~ <--- HERE
    return _1

Traceback of TorchScript, original code (most recent call last):
  File "/Users/wind/miniconda2/envs/experimental/lib/python3.7/site-packages/torchvision/models/mobilenet.py", line 190, in forward
    def forward(self, x: Tensor) -> Tensor:
        return self._forward_impl(x)
               ~~~~~~~~~~~~~~~~~~ <--- HERE
  File "/Users/wind/miniconda2/envs/experimental/lib/python3.7/site-packages/torch/nn/modules/container.py", line 117, in forward
    def forward(self, input):
        for module in self:
            input = module(input)
                    ~~~~~~ <--- HERE
        return input
  File "/Users/wind/miniconda2/envs/experimental/lib/python3.7/site-packages/torch/nn/modules/container.py", line 117, in forward
    def forward(self, input):
        for module in self:
            input = module(input)
                    ~~~~~~ <--- HERE
        return input
  File "/Users/wind/miniconda2/envs/experimental/lib/python3.7/site-packages/torch/nn/modules/conv.py", line 393, in forward
    def forward(self, input: Tensor) -> Tensor:
        return self._conv_forward(input, self.weight)
               ~~~~~~~~~~~~~~~~~~ <--- HERE
  File "/Users/wind/miniconda2/envs/experimental/lib/python3.7/site-packages/torch/nn/modules/conv.py", line 389, in _conv_forward
                            weight, self.bias, self.stride,
                            _pair(0), self.dilation, self.groups)
        return F.conv2d(input, weight, self.bias, self.stride,
               ~~~~~~~~ <--- HERE
                        self.padding, self.dilation, self.groups)
RuntimeError: expected scalar type Float but found UNKNOWN_SCALAR

Aborted`

<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->

## Expected behavior

<!-- 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): 26391143b6578c24131b652753e12d617008d0f0
 - 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. -->