diff --git a/benchmarks/dynamo/ci_expected_accuracy/inductor_torchbench_inference.csv b/benchmarks/dynamo/ci_expected_accuracy/inductor_torchbench_inference.csv index 510f3886fa4d..c62ae22a4bc8 100644 --- a/benchmarks/dynamo/ci_expected_accuracy/inductor_torchbench_inference.csv +++ b/benchmarks/dynamo/ci_expected_accuracy/inductor_torchbench_inference.csv @@ -378,7 +378,7 @@ vgg16,pass,0 -vision_maskrcnn,pass,20 +vision_maskrcnn,pass,18 diff --git a/benchmarks/dynamo/ci_expected_accuracy/inductor_torchbench_training.csv b/benchmarks/dynamo/ci_expected_accuracy/inductor_torchbench_training.csv index ff63e3c114e1..903a8aec2557 100644 --- a/benchmarks/dynamo/ci_expected_accuracy/inductor_torchbench_training.csv +++ b/benchmarks/dynamo/ci_expected_accuracy/inductor_torchbench_training.csv @@ -286,7 +286,7 @@ vgg16,pass,6 -vision_maskrcnn,pass,39 +vision_maskrcnn,pass,37 diff --git a/torch/_dynamo/variables/distributed.py b/torch/_dynamo/variables/distributed.py index 31aed55da1f2..f10d159571e2 100644 --- a/torch/_dynamo/variables/distributed.py +++ b/torch/_dynamo/variables/distributed.py @@ -274,7 +274,11 @@ class DeviceMeshVariable(DistributedVariable): if name == "get_rank": return ConstantVariable.create(self.value.get_rank()) if name == "get_local_rank": - return ConstantVariable.create(self.value.get_local_rank()) + const_args = [x.as_python_constant() for x in args] + const_kwargs = {k: v.as_python_constant() for k, v in kwargs.items()} + return ConstantVariable.create( + self.value.get_local_rank(*const_args, **const_kwargs) + ) if name == "get_group": const_args = [x.as_python_constant() for x in args] const_kwargs = {k: v.as_python_constant() for k, v in kwargs.items()}