Fix error handling and layer_type of the example

zincware · Nov 13, 2023 · 3e72dd9 · 3e72dd9
1 parent 365e691
commit 3e72dd9
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/CI/unit_tests/training_strategies/test_simple_training.py b/CI/unit_tests/training_strategies/test_simple_training.py
@@ -113,7 +113,7 @@ def test_model_error(self):
         The model is an optional input in the training strategy construction.
         The input of a model can be handled by frameworks adding the model during the
         workflow of that framework (an example is RND).
-        Testing for a KeyError if no model was added but the training method is
+        Testing for a ValueError if no model was added but the training method is
         executed.
         """
 
@@ -123,7 +123,7 @@ def test_model_error(self):
             disable_loading_bar=True,
         )
 
-        assert_raises(KeyError, trainer.train_model, self.train_ds, self.test_ds, 1)
+        assert_raises(ValueError, trainer.train_model, self.train_ds, self.test_ds, 1)
 
     def test_update_training_kwargs(self):
         """

diff --git a/examples/ResNet-Example.ipynb b/examples/ResNet-Example.ipynb
@@ -281,7 +281,8 @@
     "    embedding_size = 24, \n",
     "    hidden_sizes = [12, 12, 12], \n",
     "    depths = [2, 2, 2], \n",
-    "    layer_type = 'basic', \n",
+    "    layer_type = 'bottleneck', \n",
+    "    # layer_type = 'basic', \n",
     "    hidden_act = 'relu', \n",
     "    downsample_in_first_stage = False, \n",
     "    out_features = None, \n",