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

1

Previously searching online, and here, I didn’t see any solution yet, so just checking again.
I’m running code that works using CUDA type GPUs. I wanted to modify it to run on M1 GPUs, so I added the “mps” option. Looks like everything’s OK until it gets to the training step, where it gets the Topic’s error message. Anything close to a possible solution seems to always mention Torch RL, but I’m not using that. Any thoughts on how to fix or get around, on what better version of something to use? Details below. Thanks.

Levels:

torch               2.1.2
torchaudio          2.1.2
torchvision         0.16.2

Code/error:

import torch
from torch.utils import data
. . .

  File "/Users/gkd/Programs/CancerRiskNet/scripts/main.py", line 55, in <module>
    epoch_stats, model = train.train_model(train_data, dev_data, model, args)
  File "/Users/gkd/Programs/CancerRiskNet/cancerrisknet/learn/train.py", line 43, in train_model
    run_epoch(data_loader, train=if_train, truncate_epoch=True, models=models,
  File "/Users/gkd/Programs/CancerRiskNet/cancerrisknet/learn/train.py", line 155, in run_epoch
    batch = prepare_batch(batch, args)
  File "/Users/gkd/Programs/CancerRiskNet/cancerrisknet/learn/train.py", line 203, in prepare_batch
    batch[key] = batch[key].to(args.device)
TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64. Please use float32 instead.