How about this?
import torch                                                           
from torch import nn                                                
from torch.distribitutions.bernoulli import Bernoulli

nwords, dim = 4, 3
emb = nn.Embedding(nwords, dim)
input = torch.LongTensor([[0, 1, 2, 2, 1],
                          [0, 3, 2, 1, 2]])

out = emb(input)
rw = Bernoulli(0.3).sample((out.shape[1], ))

out[:, rw==1].mean(dim=1)