diff --git a/Project.toml b/Project.toml
index d7994b9..1f19952 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,7 +1,7 @@
 name = "NeuralEstimators"
 uuid = "38f6df31-6b4a-4144-b2af-7ace2da57606"
 authors = ["Matthew Sainsbury-Dale <msainsburydale@gmail.com> and contributors"]
-version = "0.2.2"
+version = "0.1.1"
 
 [deps]
 BSON = "fbb218c0-5317-5bc6-957e-2ee96dd4b1f0"
@@ -51,6 +51,7 @@ GraphNeuralNetworks = "0.6, 1"
 Graphs = "1"
 InvertedIndices = "1"
 Metal = "1"
+NNlib = "0.9"
 NamedArrays = "0.9, 0.10"
 NearestNeighbors = "0.4"
 RecursiveArrayTools = "2, 3"