Qualcomm: bmm over a flattened 4D input returns ~1/W of the right value when followed by unsqueeze (HTP fp16) #24502

import torch, torch.nn.functional as F
from executorch.backends.qualcomm.serialization.qc_schema import QcomChipset
from executorch.backends.qualcomm.utils.utils import (
    generate_htp_compiler_spec, generate_qnn_executorch_compiler_spec, to_edge_transform_and_lower_to_qnn)

class M(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = torch.nn.Conv2d(192, 1, 1)

    def forward(self, x):
        B, C, H, W = x.shape
        mask = F.softmax(self.conv(x).view(B, 1, H * W), -1)
        return torch.bmm(x.view(B, C, H * W), mask.transpose(1, 2)).unsqueeze(-1)

x = torch.rand(1, 192, 24, 24) * 0.5
cs = generate_qnn_executorch_compiler_spec(
    soc_model=QcomChipset.SM8850, backend_options=generate_htp_compiler_spec(use_fp16=True))
edge = to_edge_transform_and_lower_to_qnn(M().eval(), (x,), cs)