loss reduction fix 'average' to 'mean'
@@ -7,11 +7,11 @@ class KgeLoss:
def create(config):
""" Factory method for loss creation """
if config.get('train.loss') == 'ce':
- return torch.nn.CrossEntropyLoss(reduction='average')
+ return torch.nn.CrossEntropyLoss(reduction='mean')
elif config.get('train.loss') == 'bce':
- return torch.nn.BCEWithLogitsLoss(reduction='average')
+ return torch.nn.BCEWithLogitsLoss(reduction='mean')
elif config.get('train.loss') == 'kl':
- return torch.nn.KLDivLoss(reduction='average')
+ return torch.nn.KLDivLoss(reduction='mean')
else:
# perhaps TODO: try class with specified name -> extensibility
raise ValueError('train.loss')
