[mthreads]Update base/benchmarks: resolve comments of P2P_inter

FlagOpen · Dec 11, 2024 · b97e82b · b97e82b
1 parent 90b781b
commit b97e82b
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/base/benchmarks/interconnect-P2P_interserver/main.py b/base/benchmarks/interconnect-P2P_interserver/main.py
@@ -51,7 +51,10 @@ def main(config, case_config, rank, world_size, local_rank):
 
     Melements = case_config.Melements
     torchsize = (Melements, 1024, 1024)
-    tensor = torch.rand(torchsize, dtype=torch.float32).to(local_rank)
+    if "mthreads" in config.vendor:
+        tensor = torch.rand(torchsize, dtype=torch.float32).to(local_rank)
+    else:
+        tensor = torch.rand(torchsize, dtype=torch.float32).cuda()
 
     host_device_sync(config.vendor)
     multi_device_sync(config.vendor)