18 lines
661 B
Python
18 lines
661 B
Python
from ml_lib.utils.config import Config
|
|
from models.conv_classifier import ConvClassifier
|
|
from models.bandwise_conv_classifier import BandwiseConvClassifier
|
|
from models.bandwise_conv_multihead_classifier import BandwiseConvMultiheadClassifier
|
|
from models.ensemble import Ensemble
|
|
|
|
|
|
class MConfig(Config):
|
|
# TODO: There should be a way to automate this.
|
|
|
|
@property
|
|
def _model_map(self):
|
|
return dict(ConvClassifier=ConvClassifier,
|
|
BandwiseConvClassifier=BandwiseConvClassifier,
|
|
BandwiseConvMultiheadClassifier=BandwiseConvMultiheadClassifier,
|
|
Ensemble=Ensemble,
|
|
)
|