fit_a_nef.utils.get_nef

fit_a_nef.utils.get_nef(nef_cfg: Dict[str, Any]) → Module

Returns the model for the given config.

Parameters:

nef_cfg (ConfigDict) – The config for the model.

Raises:

NotImplementedError – If the model is not implemented.

Returns:

The model.

Return type:

flax.linen.Module