[D] My experience with running PyTorch on the M1 GPU

After the announcement, I was super excited to give it a try. I ran a VGG16 on both a

• an M1 MacBook Air (16 Gb RAM)
• an M1 Pro MacBook Pro (32 Gb RAM)

and the results were a bit underwhelming:

https://preview.redd.it/p8pbnptklf091.png?width=1035&format=png&auto=webp&s=26bb4a43f433b1cd983bb91c37b601b5b01c0318

The GPU performance was 2x as fast as the CPU performance on the M1 Pro, but I was hoping for more. Anyone else tried this and has any tips?

I have a more detailed write-up here: Running PyTorch on the M1 GPU

And a link to the code examples here on GitHub.