Fix Inductor runs. (#6195)

pytorch · Apr 22, 2024 · ac9e76f · ac9e76f
1 parent 9986118
commit ac9e76f
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/benchmarks/torchbench_model.py b/benchmarks/torchbench_model.py
@@ -202,8 +202,8 @@ def load_benchmark(self):
     self.module, self.example_inputs = benchmark.get_module()
 
     # Move the initialized model to XLA device.
+    device = self.benchmark_experiment.get_device()
     if self.benchmark_experiment.xla:
-      device = self.benchmark_experiment.get_device()
       self.module = self.module.to(device)
       self.example_inputs = pytree.tree_map_only(torch.Tensor,
                                                  lambda t: t.to(device),