Unfortunately, PyTorch doesn’t have docker builds for arm64, only for amd64. But you can use third-party build from another vendors.

I found builds from www.arm.com

• https://github.com/ARM-software
• https://github.com/ARM-software/Tool-Solutions/tree/main/docker/pytorch-aarch64
• https://hub.docker.com/u/armswdev
• https://hub.docker.com/r/armswdev/pytorch-arm-neoverse/tags

So you can just pull PyTorch image for arm64 from one of images https://hub.docker.com/r/armswdev/pytorch-arm-neoverse/tags