MPS device and Pytorch

Question

I want to rum mpnn with a lightning trainer on my mac. These are my trainer settings:

trainer = pl.Trainer(
    logger=False,
    enable_checkpointing=True, #
    enable_progress_bar=True,
    accelerator="mps",
    devices= 1,
    max_epochs=20, # number of epochs to train for
)

I also changed my environment varibale as followed:

import os
os.environ['PYTORCH_ENABLE_MPS_FALLBACK'] = '1'

still when I start my training:

trainer.fit(mpnn, train_loader)

I get following error message:

NotImplementedError: The operator 'aten::scatter_reduce.two_out' 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. Output is truncated. View as a scrollable element or open in a text editor. Adjust cell output settings...

I tried following steps to fix the problem:

changed setup in trainer from accelerator ='auto', to accellerator = 'mps'

changed the environment variable as suggested in the error.

tried:

if torch.backends.mps.is_available():
    device = torch.device('mps')
else:
    device = torch.device('cpu')