Update requirements.txt

Ensure Ax is a version with TensorboardMetric

Author: mpolson64

.ci/docker/requirements.txt

@@ -30,7 +30,7 @@ pytorch-lightning
 torchx
 torchrl==0.3.0
 tensordict==0.3.0
-ax-platform
+ax-platform>==0.4.0
 nbformat>==5.9.2
 datasets
 transformers