What is the synchronise function for mac mps?

mps
1

I want to let cpu to wait until the mac Neural engine tasks to finish, what function should I use? I know for CUDA, I can use torch.cuda.synchronize(), but what function is it for mps?

I have the following code:

import torch
if torch.has_mps:
    device = torch.device("mps")
else:
    device = torch.device("cpu")

print("using", device, "device")

import time

matrix_size = 32*512

x = torch.randn(matrix_size, matrix_size)
y = torch.randn(matrix_size, matrix_size)

print("***********cpu speed***************")
start = time.time()
result = torch.matmul(x,y)
print(time.time()-start)
print("verify device: ", result.device)


x_mps = x.to(device)
y_mps = y.to(device)
# torch.cuda.synchronize()


for i in range(3):
    print("***********mps speed***************")
    start = time.time()
    result_mps = torch.matmul(x_mps,y_mps)
    # torch.cuda.synchronize()
    print(time.time()-start)
    print("verify device: ", result_mps.device)