def test_parse_multiple_reg_loss(self):
loss_config = {
"regularization": [
{
"name": "bending",
"weight": 2.0,
},
{
"name": "gradient",
"weight": 2.0,
"l1": True,
},
],
}
got = parse_reg_loss(loss_config=loss_config)
assert got == loss_config
def test_parse_old_loss(self, energy_type: str, loss_name: str,
extra_args: dict):
loss_config = {
"regularization": {
"energy_type": energy_type,
"weight": 2.0,
}
}
expected = {
"regularization": {
"name": loss_name,
"weight": 2.0,
**extra_args,
},
}
got = parse_reg_loss(loss_config=loss_config)
assert got == expected