def drop_matrix_rows_3d(matrix, forget_vector):
    """
    Zeroing blocks along first dimension according to forget_vector. Forget vector should consist of 0s and 1s.

    :param matrix: size - [N1, N2, N3]
    :param forget_vector: size - [N1, 1]
    :return: size - [N1, N2, N3]
    """
    return matrix.mul(forget_vector.unsqueeze(2))