It supports PyTorch.

It also supports anything --- you give it a script and it builds a Docker container and runs it. You don't need to use any particular language or framework.

As @syllogism said, you can use any framework/library to train your models.

A couple of libraries have a special integration, Tensorflow, MXNet, Horovod, and Pytorch, they have a section in the specification for users who want to run distributed experiments.