Allow more general networks

[jwallwork23]

At the minute, we only support fully connected feed-forward networks with a single hidden layer. It would be good to build on this.

See the TODO in this commit.

[jwallwork23]

For example, perhaps we could use convolutional neural networks or graph neural networks.

[jwallwork23]

Note #14 included some small steps in this direction.