crash when running model.forward on vulkan

Hi!

I'm trying to run pytorch model on an android device and followed this tutorial :
https://github.com/pytorch/tutorials/blob/master/prototype_source/vulkan_workflow.rst
for better performance.

my problem is that when running the model using the vulkan option
mModule = Module.load(moduleFileAbsoluteFilePath, Device.VULKAN)
Tensor outputTensor = mModule.forward(IValue.from(mInputTensor)).toTensor(); crashes!
this is the crash:

2021-01-25 11:08:18.434 5280-5656/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/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/ganit/workspace/python/pytorch/tutorial-env/lib/python3.7/site-packages/torchvision/models/mobilenet.py", line 171, in forward
    def forward(self, x):
        return self._forward_impl(x)
               ~~~~~~~~~~~~~~~~~~ <--- HERE
  File "/Users/ganit/workspace/python/pytorch/tutorial-env/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/ganit/workspace/python/pytorch/tutorial-env/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/ganit/workspace/python/pytorch/tutorial-env/lib/python3.7/site-packages/torch/nn/modules/conv.py", line 423, in forward
    def forward(self, input: Tensor) -> Tensor:
        return self._conv_forward(input, self.weight)
               ~~~~~~~~~~~~~~~~~~ <--- HERE
  File "/Users/ganit/workspace/python/pytorch/tutorial-env/lib/python3.7/site-packages/torch/nn/modules/conv.py", line 419, 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

    at org.pytorch.NativePeer.forward(Native Method)
    at org.pytorch.Module.forward(Module.java:49)
    at org.pytorch.demo.vision.ImageClassificationActivity.analyzeImage(ImageClassificationActivity.java:184)

2021-01-25 11:08:18.434 5280-5656/org.pytorch.demo E/PyTorchDemo: 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:237)
at android.os.HandlerThread.run(HandlerThread.java:67)

I am running on pytorch demo.
aars comiled with vulkan as instructed in turorial.
when running mModule = Module.load(moduleFileAbsoluteFilePath, Device.VULKAN)
whithout Device.VULKAN it does not crash, but run time is very long 140ms~

Any help will be much appreciated!
Ganit

Start with prototype_source/vulkan_workflow.rst and the Android demo entry point at ImageClassificationActivity.java:184, where Module.forward is called. Reproduce the MobileNetV2 failure with Device.VULKAN, compare it with the non-Vulkan path, and trace the reported UNKNOWN_SCALAR error. Done means the Vulkan model runs without the crash and the tutorial or demo behavior is updated accordingly.