Merge branch 'main' of https://github.com/juliagmt-google/benchmark i…

…nto testing
juliagmt-google · Oct 1, 2024 · fbf2498 · fbf2498
2 parents 2e8ea36 + 252a3b1
commit fbf2498
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/torchbenchmark/operators/low_mem_dropout/operator.py b/torchbenchmark/operators/low_mem_dropout/operator.py
@@ -38,7 +38,7 @@ def triton_dropout(self, p, x):
  n_elements = x.numel()
  grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),)
 
- x_keep = (torch.rand(size=(10,)) > p).to(torch.int32).cuda()
+ x_keep = (torch.rand(size=(n_elements,)) > p).to(torch.int32).cuda()
 
  def _inner():
  return _triton_dropout[grid](