import torch


def _gen_mask(shape, drop_prob):
    """Generate a droppout mask."""
    keep_prob = 1. - drop_prob
    #mask = tf.random_uniform(shape, dtype=tf.float32)
    mask = torch.FloatTensor(shape[0], shape[1]).uniform_(0, 1)
    mask = torch.floor(mask + keep_prob) / keep_prob
    return mask