Rank of updates array (1) must be greater than or equal to inner-most dimension of indices array (29)

[wsnoble]

On an M2 Mac, I am getting the error shown below. I do have the fallback variable set properly:

% echo $PYTORCH_ENABLE_MPS_FALLBACK
1

I am using the following version:

Casanovo: 4.0.1
Depthcharge: 0.2.3
Lightning: 2.1.3
PyTorch: 2.1.2

The input MGF file is 151009_exo4_4.short.mgf.txt

% casanovo sequence 151009_exo4_4.short.mgf 
WARNING: Dataloader multiprocessing is currently not supported on Windows or MacOS; using only a single thread.
Seed set to 454
INFO: Model weights file /Users/wnoble/Library/Caches/casanovo/casanovo_massivekb_v4_0_0.ckpt retrieved from local cache
INFO: Casanovo version 4.0.1
INFO: Sequencing peptides from:
INFO:   151009_exo4_4.short.mgf
GPU available: True (mps), used: True
TPU available: False, using: 0 TPU cores
IPU available: False, using: 0 IPUs
HPU available: False, using: 0 HPUs
INFO: Reading 1 files...
151009_exo4_4.short.mgf: 2000spectra [00:00, 20073.48spectra/s]
WARNING: Skipped 1282 spectra with invalid precursor info
Predicting DataLoader 0:   0%|                            | 0/1 [00:00<?, ?it/s]WARNING: UserWarning: The operator 'aten::_nested_tensor_from_mask_left_aligned' 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:13.)
WARNING: UserWarning: The operator 'aten::_nested_tensor_from_mask_left_aligned' 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:13.)
/AppleInternal/Library/BuildRoots/0032d1ee-80fd-11ee-8227-6aecfccc70fe/Library/Caches/com.apple.xbs/Sources/MetalPerformanceShaders/MPSNDArray/Kernels/MPSNDArrayGatherND.mm:256: failed assertion `Rank of updates array (1) must be greater than or equal to inner-most dimension of indices array (29)'
zsh: abort      casanovo sequence 151009_exo4_4.short.mgf
/opt/homebrew/Caskroom/miniconda/base/envs/casanovo/lib/python3.10/multiprocessing/resource_tracker.py:224: UserWarning: resource_tracker: There appear to be 1 leaked semaphore objects to clean up at shutdown
  warnings.warn('resource_tracker: There appear to be %d '

[wsnoble]

Update: this error does not appear when I don't use the GPU.