Correct TorchXScheduler.poll_trial_status signature

Summary:
A change in core Ax's scheduler API caused this bug downstream D32762724. Currently TorchXScheduler just ignores the new option, and its default valued so this should just fix everything out of the box, but the torchx folks may or may not want to go ahead and implement the expected associated behavior in the future.

This was surfaced to me by a bisect

Reviewed By: d4l3k

Differential Revision: D32919059

fbshipit-source-id: 0fa57ed957acc090a615eb6da37ef525c88a0b0e

Author: mpolson64

torchx/runtime/hpo/ax.py

@@ -219,7 +219,9 @@ class TorchXScheduler(ax_Scheduler):
 
     """
 
-    def poll_trial_status(self) -> Dict[TrialStatus, Set[int]]:
+    def poll_trial_status(
+        self, poll_all_trial_statuses: bool = False
+    ) -> Dict[TrialStatus, Set[int]]:
         trial_statuses: Dict[TrialStatus, Set[int]] = {}
 
         for trial in self.running_trials: