Add pip dependencies needed to use tf.distribute test utils

[renning22]

Replicate commit (https://github.com/tensorflow/benchmarks/commit/3effceb5d57dcd24810ed8674b88ed5edaff4988) to cuda11 docker.