pytorch
diff --git a/‎test/inductor/test_aot_inductor.py
Lines changed: 15 additions & 4 deletions b/‎test/inductor/test_aot_inductor.py
Lines changed: 15 additions & 4 deletions
diff --git a/‎torch/csrc/inductor/aoti_runtime/model.h
Lines changed: 0 additions & 2 deletions b/‎torch/csrc/inductor/aoti_runtime/model.h
Lines changed: 0 additions & 2 deletions
@@ -1268,8 +1268,16 @@ def false_fn(x):
 
                 return torch.cond(x.shape[0] > 5, true_fn, false_fn, (x,))
 
-        input1 = (torch.ones(3, 3), torch.ones(5), torch.ones(3, 3))
-        input2 = (torch.ones(10, 3), torch.ones(6), torch.ones(10, 3))
+        input1 = (
+            torch.ones(3, 3, device=self.device),
+            torch.ones(5, device=self.device),
+            torch.ones(3, 3, device=self.device),
+        )
+        input2 = (
+            torch.ones(10, 3, device=self.device),
+            torch.ones(6, device=self.device),
+            torch.ones(10, 3, device=self.device),
+        )
         inputs = (input1, input2)
         dynamic_shapes = {"x": {0: Dim("d")}, "y": {0: Dim("d1")}, "z": {0: Dim("d")}}
         self.check_model_with_multiple_inputs(
@@ -1395,6 +1403,9 @@ def forward(self, x):
         self.check_model(M(self.device), (torch.randn(5, 5, device=self.device),))
 
     def test_zero_grid_with_backed_symbols(self):
+        if self.device != GPU_TYPE:
+            raise unittest.SkipTest("requires GPU")
+
         class Repro(torch.nn.Module):
             def __init__(self) -> None:
                 super().__init__()
@@ -1417,7 +1428,7 @@ def forward(self, x, b):
             example_inputs,
             dynamic_shapes=dynamic_shapes,
         )
-        aot_inductor_module = AOTIRunnerUtil.load(GPU_TYPE, so_path)
+        aot_inductor_module = AOTIRunnerUtil.load(self.device, so_path)
         aot_inductor_module(*example_inputs)
 
         # Re-run where dynamic dim size is 0.
@@ -1920,7 +1931,7 @@ def __init__(self) -> None:
             def forward(self, x):
                 return torch.ops.aten.normal_functional.default(x)
 
-        self.check_model(Model(), (torch.empty(4, 1, 4, 4),))
+        self.check_model(Model(), (torch.empty(4, 1, 4, 4, device=self.device),))
 
     def test_empty_graph(self):
         class Model(torch.nn.Module):
 
@@ -99,10 +99,8 @@ inline void parse_device_str(
 
   if (sm[1].str() == "cpu") {
     device_type = aoti_torch_device_type_cpu();
-#ifdef USE_CUDA
   } else if (sm[1].str() == "cuda") {
     device_type = aoti_torch_device_type_cuda();
-#endif
 #ifdef USE_XPU
   } else if (sm[1].str() == "xpu") {
     device_type = aoti_torch_device_type_xpu();