cryodrgn on Apple M1

[zhonge]

I want to generate volumes from a cryoDRGN model locally on my MacBook. I installed the nightly build of pytorch. Also had to make a few changes to loading weights and how we set the device, but it looks like there is some minor (?) refactoring we have to do in the model forward pass:

/Users/zhonge/dev/cryodrgn/cryodrgn/models.py:412: 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/_temp/anaconda/conda-bld/pytorch_1664694774655/work/aten/src/ATen/mps/MPSFallback.mm:11.)
  x = x[keep]

[zhonge]

Hmm, ran into problems in training:

  File "/Users/zhonge/dev/cryodrgn/cryodrgn/commands/train_vae.py", line 400, in main
    if in_dim % 8 != 0:
NotImplementedError: The operator 'aten::remainder.Tensor_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.