Does this work on Apple Silicon M1 GPU using mps backend?

When trying to get this working using the GPU on M1 with

Both of these return True

print(torch.backends.mps.is_available())
print(torch.backends.mps.is_built())

This also works

device = torch.device('mps')

All steps pass until this part

for x in tqdm(sampler.sample_batch_progressive(batch_size=1, model_kwargs=dict(images=[img]))):
    samples = x

Here is the error

File ~/Downloads/point-e-main/point_e/diffusion/gaussian_diffusion.py:1016, in _extract_into_tensor(arr, timesteps, broadcast_shape)
   1006 def _extract_into_tensor(arr, timesteps, broadcast_shape):
   1007     """
   1008     Extract values from a 1-D numpy array for a batch of indices.
   1009 
   (...)
   1014     :return: a tensor of shape [batch_size, 1, ...] where the shape has K dims.
   1015     """
-> 1016     res = th.from_numpy(arr).to(device=timesteps.device)[timesteps].float()
   1017     while len(res.shape) < len(broadcast_shape):
   1018         res = res[..., None]

TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64. Please use float32 instead.

openai / point-e

Does this work on Apple Silicon M1 GPU using mps backend? #28