Skip to content

[WIP] Put back MPS build on macos-m1-stable #8476

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 10 commits into from
Closed
12 changes: 3 additions & 9 deletions setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,6 @@
import shutil
import subprocess
import sys
import warnings

import torch
from pkg_resources import DistributionNotFound, get_distribution, parse_version
Expand Down Expand Up @@ -205,14 +204,9 @@ def get_extensions():
nvcc_flags = []
extra_compile_args["nvcc"] = nvcc_flags

# FIXME: MPS build breaks custom ops registration, so it was disabled.
# See https://github.com/pytorch/vision/issues/8456.
# TODO: Fix MPS build, remove warning below, and put back commented-out elif block.V
if force_mps:
warnings.warn("MPS build is temporarily disabled!!!!")
# elif torch.backends.mps.is_available() or force_mps:
# source_mps = glob.glob(os.path.join(extensions_dir, "ops", "mps", "*.mm"))
# sources += source_mps
elif torch.backends.mps.is_available() or force_mps:
source_mps = glob.glob(os.path.join(extensions_dir, "ops", "mps", "*.mm"))
sources += source_mps

if sys.platform == "win32":
define_macros += [("torchvision_EXPORTS", None)]
Expand Down
3 changes: 1 addition & 2 deletions test/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -49,8 +49,7 @@ def pytest_collection_modifyitems(items):
# There are special cases though, see below
item.add_marker(pytest.mark.skip(reason=CUDA_NOT_AVAILABLE_MSG))

# TODO: uncoment when MPS works again - see FIXME in setup.py
if needs_mps: # and not torch.backends.mps.is_available():
if needs_mps and not torch.backends.mps.is_available():
item.add_marker(pytest.mark.skip(reason=MPS_NOT_AVAILABLE_MSG))

if IN_FBCODE:
Expand Down
Loading