Unit tests are broken

It seems like the unit tests are out of date.

• test_gnn_model.py seems to be looking for a module gnn_stage that does not exist in stages.edge_classifier
• test_gnn_stage.py fails because the test_gnn_config.yaml does not have the key model