Bugs about permute function using MPS

[zhaoBowen612]

🐛 Describe the bug

class FCN(nn.Module):
def init(self):
super(FCN, self).init()
self.conv1 = nn.Conv1d(1, 128, 3)
self.gap = nn.AdaptiveAvgPool1d(1)
self.fc = nn.Linear(128, 2)

def forward(self, x):
    batch_size = x.size(0)
    x = self.conv1(x)
    x = self.gap(x)
    x = x.view(batch_size, -1)
    x = self.fc(x)
    x = torch.log_softmax(x, dim=1)
    return x

if name == 'main':
model = FCN().to('mps')
out = model(torch.randn(128, 176, 1).permute(0, 2, 1).to('mps')). # This line shows the error "Assertion failed: (mapIt != _jitValueTypes.end()), function getStaticType, file MPSRuntime_Project.h, line 435."

# out = model(torch.randn(128, 1, 176).to('mps')).  # This could work well.

Versions

PyTorch version: 1.13.0.dev20220714
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 12.4 (arm64)
GCC version: Could not collect
Clang version: 13.1.6 (clang-1316.0.21.2.5)
CMake version: version 3.22.1
Libc version: N/A

Python version: 3.9.12 | packaged by conda-forge | (main, Mar 24 2022, 23:25:14) [Clang 12.0.1 ] (64-bit runtime)
Python platform: macOS-12.4-arm64-arm-64bit
Is CUDA available: False
CUDA runtime version: No CUDA
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] numpy==1.21.5
[pip3] torch==1.13.0.dev20220714
[pip3] torchaudio==0.14.0.dev20220603
[pip3] torchvision==0.14.0.dev20220714
[conda] torch 1.13.0.dev20220714 pypi_0 pypi
[conda] torchaudio 0.14.0.dev20220603 pypi_0 pypi
[conda] torchvision 0.14.0.dev20220714 pypi_0 pypi

cc @kulinseth @albanD

[qqaatw]

I can reproduce. The repro can be simplified as the following:

def repro():
    model = torch.nn.Conv1d(1, 128, 3)
    a = torch.ones(128, 176, 1).permute(0, 2, 1)
    out = model(a)  # pass

    a_mps = a.to("mps")
    model = model.to("mps")
    out = model(a_mps)  # fail

[DenisVieriu97]

Thanks @zhaoBowen612 and @qqaatw for the repro. We'll take a look at the crash

[zhaoBowen612]

It seems that the function "transpose" has the same problem...

[zhaoBowen612]

It seems the problem still exists? @DenisVieriu97

[DenisVieriu97]

@zhaoBowen612 the fix has not been merged in PyTorch master yet. I'll update this thread when the change gets merged