diff --git a/CHANGELOG.md b/CHANGELOG.md
index a7895ec3..308197e3 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -3,9 +3,11 @@
 ## BUG FIXES
 
 * Fix the component paths, build paths and `rename_keys` separator in the helper scripts, which prevented `scripts/create_datasets/test_resources.sh` and both `run_test.sh` scripts from running at all (PR #22).
-* 
+
 * `cellmapper_scvi`: Fix postfix due to breaking changes in package (PR #23).
 
+* `simple_mlp_predict`: Size the model output with `input_train_mod2.n_vars` instead of `input_test_mod1.n_vars`. The two only coincide on the test resources, so this crashed on every real dataset (PR #24).
+
 # task_predict_modality 0.1.1
 
 ## NEW FUNCTIONALITY
diff --git a/src/methods/simple_mlp/simple_mlp_predict/script.py b/src/methods/simple_mlp/simple_mlp_predict/script.py
index 2509c8ac..aaf31647 100644
--- a/src/methods/simple_mlp/simple_mlp_predict/script.py
+++ b/src/methods/simple_mlp/simple_mlp_predict/script.py
@@ -61,7 +61,7 @@ def _predict(model,dl):
 
 print('Start predict', flush=True)
 if task == 'GEX2ATAC':
-    y_pred = ymean*np.ones([input_test_mod1.n_obs, input_test_mod1.n_vars])
+    y_pred = ymean*np.ones([input_test_mod1.n_obs, input_train_mod2.n_vars])
 else:
     folds = [0, 1, 2]
 
@@ -82,7 +82,7 @@ def _predict(model,dl):
         model_inf = MLP.load_from_checkpoint(
             ckpt,
             in_dim=X.shape[1],
-            out_dim=input_test_mod1.n_vars,
+            out_dim=input_train_mod2.n_vars,
             ymean=ymean,
             config=config
         )