Repository: pytorch
Issue: Segmentation Faults loading Models pytorch v2.3.0 Apple M2 
Labels: high priority,module: crash,triaged,module: macos,module: openmp

### 🐛 Describe the bug

Segementation faults loading a UNet model on pytorch v2.3.0 on macos Apple M2.

likely not a UNet specific things but its the quickest model I have at hand to easily reproduce this. Minimum reproducible examples in the pytest tests below. 

The last version where these models seem loadable is pytorch v2.1.2. 

```python
from pathlib import Path
from typing import cast

import torch

from torch import Tensor, nn

# region private


class _DoubleConv(nn.Module):
    def __init__(self, in_channels: int, out_channels: int) -> None:
        super().__init__()
        self.double_conv = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
        )

    def forward(self, x: Tensor) -> Tensor:
        return cast(Tensor, self.double_conv(x))


class _Down(nn.Module):
    def __init__(self, in_channels: int, out_channels: int) -> None:
        super().__init__()
        self.maxpool_conv = nn.Sequential(
            nn.MaxPool2d(2), _DoubleConv(in_channels, out_channels)
        )

    def forward(self, x: Tensor) -> Tensor:
        return cast(Tensor, self.maxpool_conv(x))


class _Up(nn.Module):
    def __init__(self, in_channels: int, out_channels: int) -> None:
        super().__init__()
        self.up = nn.Upsample(scale_factor=2, mode="bilinear", align_corners=True)
        self.conv = _DoubleConv(in_channels, out_channels)

    def forward(self, x1: Tensor, x2: Tensor) -> Tensor:
        x1 = self.up(x1)
        x = torch.cat([x2, x1], dim=1)
        return cast(Tensor, self.conv(x))


class _OutConv(nn.Module):
    def __init__(self, in_channels: int, out_channels: int) -> None:
        super().__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1)

    def forward(self, x: Tensor) -> Tensor:
        return cast(Tensor, self.conv(x))


# endregion

# region model


class UNet(nn.Module):
    def __init__(self, n_channels: int, n_classes: int) -> None:
        super().__init__()
        self.n_channels = n_channels
        self.n_classes = n_classes
        self.inc = _DoubleConv(n_channels, 64)
        self.down1 = _Down(64, 128)
        self.down2 = _Down(128, 256)
        self.down3 = _Down(256, 512)
        self.down4 = _Down(512, 512)
        self.up1 = _Up(1024, 256)
        self.up2 = _Up(512, 128)
        self.up3 = _Up(256, 64)
        self.up4 = _Up(128, 64)
        self.outc = _OutConv(64, n_classes)

    def forward(self, x: Tensor) -> Tensor:
        x1 = self.inc(x)
        x2 = self.down1(x1)
        x3 = self.down2(x2)
        x4 = self.down3(x3)
        x5 = self.down4(x4)
        x = self.up1(x5, x4)
        x = self.up2(x, x3)
        x = self.up3(x, x2)
        x = self.up4(x, x1)
        logits = self.outc(x)
        return cast(Tensor, logits)


class TestModelLoading:
    def test_load_model__using_filelikes(
        self,
        tmp_path: Path,
    ) -> None:
        model_path = tmp_path / "model.pkl"

        with model_path.open("wb") as f:
            torch.save({"model_state": UNet(n_channels=3, n_classes=2).state_dict()}, f)

        with model_path.open("rb") as f:
            state_dict = torch.load(f, map_location=torch.device("cpu"))["model_state"]
            UNet(n_channels=3, n_classes=2).load_state_dict(state_dict)

    def test_load_model__using_filepaths(self, tmp_path: Path) -> None:
        model_path = tmp_path / "model.pkl"

        torch.save(
            {"model_state": UNet(n_channels=3, n_classes=2).state_dict()},
            model_path,
        )

        state_dict = torch.load(model_path, map_location=torch.device("cpu"))[
            "model_state"
        ]
        UNet(n_channels=3, n_classes=2).load_state_dict(state_dict)

```

Resulting error for both of these is:

```bash
Thread 0x00000001e876fac0 (most recent call first):
  File "xxxx/python3.12/site-packages/torch/nn/modules/module.py", line 2075 in _load_from_state_dict
  File "xxxx/python3.12/site-packages/torch/nn/modules/module.py", line 2157 in load
  File "xxxx/python3.12/site-packages/torch/nn/modules/module.py", line 2163 in load
/var/folders/38/7t22pn0n7zj57vk572qkksgc0000gn/T/tmpa0gw1xzz: line 3: 34318 Segmentation fault: 11 
```

### Versions

PyTorch version: 2.3.0
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 14.4.1 (arm64)
GCC version: Could not collect
Clang version: 16.0.6
CMake version: Could not collect
Libc version: N/A

Python version: 3.12.3 | packaged by conda-forge | (main, Apr 15 2024, 18:35:20) [Clang 16.0.6 ] (64-bit runtime)
Python platform: macOS-14.4.1-arm64-arm-64bit
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
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

CPU:
Apple M2 Pro

Versions of relevant libraries:
[pip3] mypy==1.10.0
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.26.4
[pip3] pytorch-lightning==2.2.2
[pip3] torch==2.3.0
[pip3] torchmetrics==1.3.2
[conda] libtorch                  2.3.0           cpu_generic_hed86ac2_0    conda-forge
[conda] nomkl                     1.0                  h5ca1d4c_0    conda-forge
[conda] numpy                     1.26.4          py312h8442bc7_0    conda-forge
[conda] pytorch                   2.3.0           cpu_generic_py312h16b77fb_0    conda-forge
[conda] pytorch-lightning         2.2.2              pyhd8ed1ab_0    conda-forge
[conda] torchmetrics              1.3.2              pyhd8ed1ab_0    conda-forge

cc @ezyang @gchanan @zou3519 @kadeng @msaroufim @malfet @albanD