GPU usage in Apple Silicon

[medsouza]

Hi there,

I would like to know how to ensure GPU usage in MacBook equipped with GPU (i.e. torch.backends.mps.is_available() returns True) running DeepXDE with PyTorch (and Tensorflow).

Thank you very much in advance.

Cheers,
M

[medsouza]

I experimented with the following and it worked. I don't know if this is the smartest way to go about it.

import deepxde as dde
from deepxde.backend import torch

if torch.backends.mps.is_available():
torch.set_default_device("mps")