requirements

models/ensemble.py

@@ -8,17 +8,17 @@ from torch.nn import ModuleList
 from ml_lib.modules.utils import LightningBaseModule
 from ml_lib.utils.config import Config
 from ml_lib.utils.model_io import SavedLightningModels
-from util.module_mixins import (BaseOptimizerMixin, BaseTrainMixin, BaseValMixin, BinaryMaskDatasetFunction,
+from util.module_mixins import (BaseOptimizerMixin, BaseTrainMixin, BaseValMixin, BinaryMaskDatasetMixin,
                                 BaseDataloadersMixin)
 
 
-class Ensemble(BinaryMaskDatasetFunction,
-                     BaseDataloadersMixin,
-                     BaseTrainMixin,
-                     BaseValMixin,
-                     BaseOptimizerMixin,
-                     LightningBaseModule
-                     ):
+class Ensemble(BinaryMaskDatasetMixin,
+               BaseDataloadersMixin,
+               BaseTrainMixin,
+               BaseValMixin,
+               BaseOptimizerMixin,
+               LightningBaseModule
+               ):
 
     def __init__(self, hparams):
         super(Ensemble, self).__init__(hparams)