def get_criterion(self, stage: str) -> _Criterion:
criterion_params = \
self.stages_config[stage].get("criterion_params", {})
criterion = CRITERIONS.get_from_params(**criterion_params)
if torch.cuda.is_available():
criterion = criterion.cuda()
return criterion
def _get_criterion(**params):
key_value_flag = params.pop("_key_value", False)
if key_value_flag:
criterion = {}
for key, params_ in params.items():
criterion[key] = ConfigExperiment._get_criterion(**params_)
else:
criterion = CRITERIONS.get_from_params(**params)
if criterion is not None and torch.cuda.is_available():
criterion = criterion.cuda()
return criterion