import torch


def get_class_weights(dataset, indices, verbose=False):
    """Returns class weights of a subset defined by indices
    """

    # get labels in subset
    labels = [dataset.targets[i] for i in indices]
    class_weights = 1 / torch.Tensor([labels.count(0), labels.count(1), labels.count(2),])

    if verbose:
        print(f"class weights are: {str(class_weights)}")

    return class_weights