Make sure the ResNet18 model is included in the list of available mod…

…els.
easeml · Mar 1, 2024 · 542684b · 542684b
1 parent 52742e5
commit 542684b
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/experiments/datascope/experiments/scenarios/datascope_scenario.py b/experiments/datascope/experiments/scenarios/datascope_scenario.py
@@ -35,6 +35,7 @@ class ModelSpec(str, Enum):
     NaiveBayes = "nb"
     NeuralNetwork = "nn"
     XGBoost = "xgb"
+    ResNet18 = "resnet-18"
 
 
 MODEL_TYPES = {
@@ -52,6 +53,7 @@ class ModelSpec(str, Enum):
     ModelSpec.NaiveBayes: ModelType.NaiveBayes,
     ModelSpec.NeuralNetwork: ModelType.NeuralNetwork,
     ModelSpec.XGBoost: ModelType.XGBoost,
+    ModelSpec.ResNet18: ModelType.ResNet18,
 }
 
 MODEL_KWARGS: Dict[ModelSpec, Dict[str, Any]] = {
@@ -69,6 +71,7 @@ class ModelSpec(str, Enum):
     ModelSpec.NaiveBayes: {},
     ModelSpec.NeuralNetwork: {},
     ModelSpec.XGBoost: {},
+    ModelSpec.ResNet18: {},
 }