Set correct cuda.current_device for multi-device onnx performance ben…

…ch (#115670)

Summary:
Otherwise `torch.cuda.synchronize()` works on a different device from the one that
runs PyTorch model, which lead to incorrect performance number.

X-link: pytorch/pytorch#115670
Approved by: https://github.com/thiagocrepaldi

Reviewed By: jeanschmidt

Differential Revision: D52244270

fbshipit-source-id: 3b4133041ceec63aa6f5b72b840434d468345da7

userbenchmark/dynamo/dynamobench/common.py

@@ -789,6 +789,11 @@ def timed_onnx(model, onnx_model: OnnxModel, inputs):
             if should_randomize_input
             else example_inputs
         )
+        if torch.cuda.device_count() > 1:
+            # Manually set correct torch.cuda.current_device to ensure torch.cuda.synchronize() works as intended.
+            # When there are more than 1 cuda devices, the first one is used for pytorch eager.
+            # The second one is used for onnx ort.
+            torch.cuda.set_device(0)
         timings[rep, 0], expected_output = timed(
             model,
             model_iter_fn,
@@ -797,7 +802,11 @@ def timed_onnx(model, onnx_model: OnnxModel, inputs):
             times=times,
             collect_outputs=args.collect_outputs,
         )
-
+        if torch.cuda.device_count() > 1:
+            # Manually set correct torch.cuda.current_device to ensure torch.cuda.synchronize() works as intended.
+            # When there are more than 1 cuda devices, the first one is used for pytorch eager.
+            # The second one is used for onnx ort.
+            torch.cuda.set_device(1)
         timings[rep, 1], actual_output = timed_onnx(model, onnx_model, inputs)
 
     pvalue = ttest_ind(timings[:, 0], timings[:, 1]).pvalue