diff --git a/examples/demo_dlmbl/python/excercise_1.py b/examples/demo_dlmbl/python/excercise_1.py
index b7bf6138..f2af135c 100644
--- a/examples/demo_dlmbl/python/excercise_1.py
+++ b/examples/demo_dlmbl/python/excercise_1.py
@@ -177,6 +177,8 @@
 # %%
 model_config = {
     "architecture": "2D",
+    "in_channels": 1,
+    "out_channels": 2,
     "residual": True,
     "dropout": 0.1,
     "task": "reg",