ns-train splatfacto fails to run on Apple Silicon, nerfacto works ok

[Itox001]

Describe the bug
I have an M3 and could successfully install nerfstudio, and (albeit slowly) train my first model using nerfacto. The problem is that when I try to use splatfacto, it exceptions out, saying:

File "/Users/user/dev/oss/nerfstudio/nerfstudio/models/splatfacto.py", line 229, in populate_modules
shs = torch.zeros((self.seed_points[1].shape[0], dim_sh, 3)).float().cuda()
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/homebrew/Caskroom/miniconda/base/envs/nerfstudio/lib/python3.11/site-packages/torch/cuda/init.py", line 289, in _lazy_init
raise AssertionError("Torch not compiled with CUDA enabled")
AssertionError: Torch not compiled with CUDA enabled

Naturally my torch does not have CUDA enabled since there is no Nvidia GPU. I tried with --machine.device-type=cpu, with the same result.

Command that works:
ns-train nerfacto --machine.device-type=mps --max-num-iterations 2000 --data /Users/user/dev/oss/nerfstudio/data/nerfstudio/desk

Command that produces the above exception:
ns-train splatfacto --machine.device-type=mps --max-num-iterations 2000 --data /Users/user/dev/oss/nerfstudio/data/nerfstudio/desk

Additional context
nerfstudio 1.1.3
torch 2.1.2
torchvision 0.16.2

Thanks!

[maturk]

Splatfacto requires CUDA for training right now. If you are very keen and willing to hack away, it is in theory possible to train with cpu torch right now using the _torch_impl from the cuda backend but its not supported out of the box.

[Itox001]

Thanks for the quick response @maturk ! I see, support for MPS across the ecosystem seems to be evolving quickly now, there is a lot of active development, but there is still a lot of ground to cover.

One last question realted to this topic. There are several improvements for MPS on the latest PyTorch versions, but when I tried to upgrade it to latest, I could not longer run ns-train. Do you think it is hard to resolve those conflicts? The error thrown is a bit cryptic:

[1] 27951 segmentation fault PYTORCH_ENABLE_MPS_FALLBACK=1 ns-train nerfacto --machine.device-type=mps

Without the fallback flag:

OMP: Error #179: Function pthread_mutex_init failed:
OMP: System error #22: Invalid argument
[1] 28347 abort ns-train nerfacto --machine.device-type=cpu --max-num-iterations 2000 --data

[Itox001]

Original question answered, no need to keep open.