Not able to train on apple silicon gpu "mps"

[kzaheer09]

You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.
/Users/zaheer/MinImagen/minimagen/training.py:88: UserWarning: The dst MTL buffer in copy_to_mps is non-contiguous (Triggered internally at /Users/runner/work/pytorch/pytorch/pytorch/aten/src/ATen/native/mps/operations/Copy.mm:358.)
batch[didx][tensor] = batch[didx][tensor].to(self.device)
loc("mps_subtract"("(mpsFileLoc): /AppleInternal/Library/BuildRoots/9e200cfa-7d96-11ed-886f-a23c4f261b56/Library/Caches/com.apple.xbs/Sources/MetalPerformanceShadersGraph/mpsgraph/MetalPerformanceShadersGraph/Core/Files/MPSGraphUtilities.mm":228:0)): error: input types 'tensor<64x1xf32>' and 'tensor<64x8xsi64>' are not broadcast compatible
LLVM ERROR: Failed to infer result type(s).
zsh: abort python train.py
(base) zaheer$: /Users/zaheer/.pyenv/versions/3.10.6/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 '