import torch


def drop_connect(inputs: torch.Tensor, drop_connect_rate: float, training: bool):
    """ Randomly mask a set of samples. Provides similar regularization as stochastic depth.

    Args:
        input (torch.Tensor): Input tensor to mask.
        drop_connect_rate (float): Probability of droppping each sample.
        training (bool): Whether or not the model is training
    """

    if not training:
        return inputs

    keep_prob = 1 - drop_connect_rate
    rand_tensor = keep_prob + torch.rand(
        [inputs.size()[0], 1, 1, 1],
        dtype=inputs.dtype,
        device=inputs.device,
    )
    rand_tensor.floor_()  # binarize
    output = inputs.div(keep_prob) * rand_tensor
    return output