Count labels using collection package

count_labels.py

import collection
log.debug(len(train_loader))
labels = []
for i, data in enumerate(train_loader):
    _, label = data
    labels.append(label.item())
labels_counter = collections.Counter(labels)
log.debug(labels_counter)
labels = []
log.debug(len(val_loader))
for i, data in enumerate(val_loader):
    _, label = data
    labels.append(label.item())
labels_counter = collections.Counter(labels)
log.debug(labels_counter)