How to use the m1 chip macbook to use the gpu acceleration on d2l?

Question

d2l: https://d2l.ai/

I modified def try_gpu(i=0): and def try_all_gpus() in the torch.py ​​:

def try_gpu(i=0): """Return gpu(i) if exists, otherwise return cpu().

Defined in :numref:`sec_use_gpu`"""
if torch.cuda.device_count() >= i + 1:
    return torch.device(f'cuda:{i}')
try:
    return torch.device('mps')
except:
    return torch.device('cpu')

def try_all_gpus(): """Return all available GPUs, or [cpu(),] if no GPU exists.

Defined in :numref:`sec_use_gpu`"""
devices = [torch.device(f'cuda:{i}')
           for i in range(torch.cuda.device_count())]
try:
    device_macos = torch.device('mps')
except:
    device_macos = torch.device('cpu')
return devices if devices else [device_macos]