keras

keras copied to clipboard

First off, congratulations on keras-core: keras is awesome, keras-core is awesomer! Using a Mac, I was trying to manually set a keras-core more with torch backend to benefit from the Metal GPU acceleration, which works on both Apple silicon and AMD GPUs. I noted this was not possible due to:

• get_device(), which reads from a threading global state.
• the default device is CUDA

As can be seen here, in MacOS systems using the GPU can be as easy as setting the device to "mps". So I would like to ask if there would be an opportunity for a PR that implements some version of the snippet below that wouldn't fail on non-Apple machines:

# keras_core/backend/torch/core.py
if torch.cuda.is_available():
    DEFAULT_DEVICE = "cuda"
elif:
    if torch.backends.mps.is_available():
    DEFAULT_DEVICE = "mps"
else:
    DEFAULT_DEVICE = "cpu"

The goal is to have Keras-core seamlessly choose MPS as the default GPU in PyTorch backends, just as it already does for tensorflow implementations using the Metal pluggable device.