Fix hardcoded shape in low_mem_dropout benchmark (#2475)

Summary: Pull Request resolved: #2475

Reviewed By: htyu

Differential Revision: D63653081

Pulled By: xuzhao9

fbshipit-source-id: 8d840986779b6124cbccc2425c24e2b892d55ce4

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](