# These requirements are used for the CI and CPU-only Docker images so we install CPU only versions of torch.
# For GPU Docker images, you should install requirements_ml_docker.txt afterwards.

tensorflow==2.6.2
tensorflow-probability==0.14.1

# If you make changes to the torch versions below, please also make the corresponding changes to `requirements_ml_docker.txt`!

torch==1.9.0;sys_platform=="darwin"
torchvision==0.10.0;sys_platform=="darwin"

# On non-OSX machines only install CPU version of torch and torchvision
-f https://download.pytorch.org/whl/torch_stable.html
torch==1.9.0+cpu;sys_platform!="darwin"
torchvision==0.10.0+cpu;sys_platform!="darwin"