Fix walrus operator precedence. (#6374)

pytorch · Jan 25, 2024 · 44660d8 · 44660d8
1 parent 4ab7a24
commit 44660d8
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/benchmarks/torchbench_model.py b/benchmarks/torchbench_model.py
@@ -300,7 +300,7 @@ def load_benchmark(self):
     # torch.backends.__allow_nonbracketed_mutation_flag = True
 
     # torchbench uses `xla` as device instead of `tpu`
-    if device := self.benchmark_experiment.accelerator == 'tpu':
+    if (device := self.benchmark_experiment.accelerator) == 'tpu':
       device = str(self.benchmark_experiment.get_device())
     return benchmark_cls(
         test=self.benchmark_experiment.test,