ptrblck:

torch.utils.benchmark


Thanks for your reply. I tried torch.utils.benchmark but I got similar results
import torch
n=100000
dim=2000
device='cuda'
import torch.utils.benchmark as benchmark

# masked_scatter_
x = torch.rand(n, dim, device=device)
y =  torch.rand(n, dim, device=device)
condition = torch.rand(n,1, device=device) > 0.5
print(benchmark.Timer('x.masked_scatter_(condition.expand(n,dim), y)',globals={'x':x,'y':y,'condition':condition,'n':n,'dim':dim}).timeit(10))

# torch.where
x = torch.rand(n, dim, device=device)
y =  torch.rand(n, dim, device=device)
condition = torch.rand(n,1, device=device) > 0.5
print(benchmark.Timer('torch.where(condition.expand(n,dim), x, y)',globals={'x':x,'y':y,'condition':condition,'n':n,'dim':dim}).timeit(10))

Output:
<torch.utils.benchmark.utils.common.Measurement object at 0x0000019B0E340670>
x.masked_scatter_(condition.expand(n,dim), y)
  22.45 ms
  1 measurement, 10 runs , 1 thread
<torch.utils.benchmark.utils.common.Measurement object at 0x0000019B097ADCA0>
torch.where(condition.expand(n,dim), x, y)
  5.53 ms
  1 measurement, 10 runs , 1 thread