update: torch for Linux on ARM isn't built with Apple's MPS support so it didn't work with the pip install version. Perhaps it's possible to compile from scratch to have it.