「...'aten::nonzero' is not currently supported...」on M1

[changchiyou]

Problem:
muti error occurred when I try to run the example from help(MTCNN) on my MacBook Air M1.
Env:
macOS Monterey 12.6
MacBook Air(M1, 2020) 8GB RAM
iBoot 7459.141.1

conda 4.13.0
Python 3.8.13
# packages in environment at /opt/homebrew/anaconda3/envs/pytorch:
#
# Name Version Build Channel
bzip2 1.0.8 h3422bc3_4 conda-forge
ca-certificates 2022.9.24 h4653dfc_0 conda-forge
certifi 2022.9.24 pypi_0 pypi
charset-normalizer 2.1.1 pypi_0 pypi
facenet-pytorch 2.5.2 pypi_0 pypi
idna 3.4 pypi_0 pypi
libffi 3.4.2 h3422bc3_5 conda-forge
libsqlite 3.39.4 h76d750c_0 conda-forge
libzlib 1.2.12 h03a7124_4 conda-forge
ncurses 6.3 h07bb92c_1 conda-forge
numpy 1.23.4 pypi_0 pypi
opencv-python 4.6.0.66 pypi_0 pypi
openssl 3.0.5 h03a7124_2 conda-forge
pillow 9.2.0 pypi_0 pypi
pip 22.2.2 pyhd8ed1ab_0 conda-forge
python 3.8.13 hd3575e6_0_cpython conda-forge
readline 8.1.2 h46ed386_0 conda-forge
requests 2.28.1 pypi_0 pypi
setuptools 65.4.1 pyhd8ed1ab_0 conda-forge
sqlite 3.39.4 h2229b38_0 conda-forge
tk 8.6.12 he1e0b03_0 conda-forge
torch 1.14.0.dev20221012 pypi_0 pypi
torchvision 0.15.0.dev20221012 pypi_0 pypi
typing-extensions 4.4.0 pypi_0 pypi
urllib3 1.26.12 pypi_0 pypi
wheel 0.37.1 pyhd8ed1ab_0 conda-forge
xz 5.2.6 h57fd34a_0 conda-forge

Result:
(pytorch) Running on device: mps /opt/homebrew/anaconda3/envs/pytorch/lib/python3.8/site-packages/facenet_pytorch/models/utils/detect_face.py:210: UserWarning: The operator 'aten::nonzero' is not currently supported on the MPS backend and will fall back to run on the CPU. This may have performance implications. (Triggered internally at /Users/runner/work/pytorch/pytorch/pytorch/aten/src/ATen/mps/MPSFallback.mm:11.) mask_inds = mask.nonzero() Traceback (most recent call last): File "facenetMTCNN_example.py", line 13, in <module> boxes, probs, points = mtcnn.detect(img, landmarks=True) File "/opt/homebrew/anaconda3/envs/pytorch/lib/python3.8/site-packages/facenet_pytorch/models/mtcnn.py", line 313, in detect batch_boxes, batch_points = detect_face( File "/opt/homebrew/anaconda3/envs/pytorch/lib/python3.8/site-packages/facenet_pytorch/models/utils/detect_face.py", line 79, in detect_face pick = batched_nms(boxes_scale[:, :4], boxes_scale[:, 4], image_inds_scale, 0.5) File "/opt/homebrew/anaconda3/envs/pytorch/lib/python3.8/site-packages/torchvision/ops/boxes.py", line 75, in batched_nms return _batched_nms_coordinate_trick(boxes, scores, idxs, iou_threshold) File "/opt/homebrew/anaconda3/envs/pytorch/lib/python3.8/site-packages/torch/jit/_trace.py", line 1136, in wrapper return fn(*args, **kwargs) File "/opt/homebrew/anaconda3/envs/pytorch/lib/python3.8/site-packages/torchvision/ops/boxes.py", line 94, in _batched_nms_coordinate_trick keep = nms(boxes_for_nms, scores, iou_threshold) File "/opt/homebrew/anaconda3/envs/pytorch/lib/python3.8/site-packages/torchvision/ops/boxes.py", line 41, in nms return torch.ops.torchvision.nms(boxes, scores, iou_threshold) File "/opt/homebrew/anaconda3/envs/pytorch/lib/python3.8/site-packages/torch/_ops.py", line 442, in __call__ return self._op(*args, **kwargs or {}) NotImplementedError: The operator 'torchvision::nms' is not currently implemented for the MPS device. If you want this op to be added in priority during the prototype phase of this feature, please comment on https://github.com/pytorch/pytorch/issues/77764. As a temporary fix, you can set the environment variable `PYTORCH_ENABLE_MPS_FALLBACK=1` to use the CPU as a fallback for this op. WARNING: this will be slower than running natively on MPS.

Here's the code:
from PIL import Image, ImageDraw
from facenet_pytorch import MTCNN, extract_face
import cv2
import torch

device = 'mps' if torch.backends.mps.is_available() and torch.backends.mps.is_built() else 'cpu'
print("Running on device: {}".format(device))

#cap = cv2.VideoCapture(0)
img = Image.open('./photos/Chris/16871.jpg')

mtcnn = MTCNN(keep_all=True, device=device)
boxes, probs, points = mtcnn.detect(img, landmarks=True)
print(boxes, probs, points)

img_draw = img.copy()
draw = ImageDraw.Draw(img_draw)
for i, (box, point) in enumerate(zip(boxes, points)):
    draw.rectangle(box.tolist(), width=5)
    for p in point:
        draw.rectangle((p - 10).tolist() + (p + 10).tolist(), width=10)
    extract_face(img, box, save_path='detected_face_{}.png'.format(i))
img_draw.save('annotated_faces.png')