[Tp Test] Fixe the placment of the device tensor (#1054)

pytorch · Oct 11, 2024 · 5277507 · 5277507
1 parent ec860a1
commit 5277507
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/torchao/testing/utils.py b/torchao/testing/utils.py
@@ -309,6 +309,8 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         y_q = dn_quant(up_quant(example_input))
 
         mesh = self.build_device_mesh()
+        mesh.device_type = "cuda"
+
         # Shard the models
         up_dist = self.colwise_shard(up_quant, mesh)
         dn_dist = self.rowwise_shard(dn_quant, mesh)