diff --git a/poetry.lock b/poetry.lock
index dd1bae0..33355ea 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -2199,6 +2199,32 @@ pyopenssl = ["cryptography (>=38.0.3)", "pyopenssl (>=20.0.0)"]
 reauth = ["pyu2f (>=0.1.5)"]
 requests = ["requests (>=2.20.0,<3.0.0.dev0)"]
 
+[[package]]
+name = "gpytorch"
+version = "1.13"
+description = "An implementation of Gaussian Processes in Pytorch"
+optional = false
+python-versions = ">=3.8"
+files = [
+    {file = "gpytorch-1.13-py3-none-any.whl", hash = "sha256:97da5b07a524952612e8d265ec89d4a5ec0ef0587e76d6178961ce26ef9679d1"},
+    {file = "gpytorch-1.13.tar.gz", hash = "sha256:f4a488633a2a7a4ab37d12553d1d1dd39690043dbceef14ca428b7d5f89f73ba"},
+]
+
+[package.dependencies]
+jaxtyping = "0.2.19"
+linear-operator = ">=0.5.3"
+mpmath = ">=0.19,<=1.3"
+scikit-learn = "*"
+scipy = ">=1.6.0"
+
+[package.extras]
+dev = ["pre-commit", "setuptools-scm", "twine", "ufmt"]
+docs = ["ipykernel (<=6.17.1)", "ipython (<=8.6.0)", "lxml-html-clean", "m2r2 (<=0.3.3.post2)", "nbclient (<=0.7.3)", "nbformat (<=5.8.0)", "nbsphinx (<=0.9.1)", "platformdirs (<=3.2.0)", "setuptools-scm (<=7.1.0)", "sphinx (<=6.2.1)", "sphinx-autodoc-typehints (<=1.23.0)", "sphinx-rtd-theme (<0.5)", "uncompyle6 (<=3.9.0)"]
+examples = ["ipython", "jupyter", "matplotlib", "scipy", "torchvision", "tqdm"]
+keops = ["pykeops (>=1.1.1)"]
+pyro = ["pyro-ppl (>=1.8)"]
+test = ["flake8 (==4.0.1)", "flake8-print (==4.0.0)", "nbval", "pytest"]
+
 [[package]]
 name = "graphene"
 version = "3.3"
@@ -2834,6 +2860,22 @@ files = [
     {file = "jabbar-0.0.16.tar.gz", hash = "sha256:522f29ca04e44a25fbc3ae0419f7d7bf96af1b9d131bbd8b58899224fc5eb0f5"},
 ]
 
+[[package]]
+name = "jaxtyping"
+version = "0.2.19"
+description = "Type annotations and runtime checking for shape and dtype of JAX arrays, and PyTrees."
+optional = false
+python-versions = "~=3.8"
+files = [
+    {file = "jaxtyping-0.2.19-py3-none-any.whl", hash = "sha256:651352032799d422987e783fd1b77699b53c3bb28ffa644bbca5f75ec4fbb843"},
+    {file = "jaxtyping-0.2.19.tar.gz", hash = "sha256:21ff4c3caec6781cadfe980b019dde856c1011e17d11dfe8589298040056325a"},
+]
+
+[package.dependencies]
+numpy = ">=1.20.0"
+typeguard = ">=2.13.3"
+typing-extensions = ">=3.7.4.1"
+
 [[package]]
 name = "jedi"
 version = "0.19.1"
@@ -3269,6 +3311,28 @@ cli = ["fire"]
 docs = ["requests (>=2.0.0)"]
 typing = ["mypy (>=1.0.0)", "types-setuptools"]
 
+[[package]]
+name = "linear-operator"
+version = "0.5.3"
+description = "A linear operator implementation, primarily designed for finite-dimensional positive definite operators (i.e. kernel matrices)."
+optional = false
+python-versions = ">=3.8"
+files = [
+    {file = "linear_operator-0.5.3-py3-none-any.whl", hash = "sha256:908df4e64e25312edfa5502b30b71df97383cb604a13f420921030ae40c47838"},
+    {file = "linear_operator-0.5.3.tar.gz", hash = "sha256:16122661cd8b8a89ea965c845f650affe0f688f315893bb8dfa1182f148a1a41"},
+]
+
+[package.dependencies]
+jaxtyping = "0.2.19"
+mpmath = ">=0.19,<=1.3"
+scipy = "*"
+torch = ">=1.11"
+
+[package.extras]
+dev = ["pre-commit", "setuptools-scm", "twine", "ufmt"]
+docs = ["myst-parser", "setuptools-scm", "six", "sphinx", "sphinx-autodoc-typehints", "sphinx-rtd-theme", "uncompyle6 (<=3.9.0)"]
+test = ["flake8 (==5.0.4)", "flake8-print (==5.0.0)", "pytest", "typeguard (>=2.13.3,<2.14.0)"]
+
 [[package]]
 name = "llvmlite"
 version = "0.43.0"
@@ -8613,4 +8677,4 @@ dev = ["pre-commit", "pytest", "ruff"]
 [metadata]
 lock-version = "2.0"
 python-versions = ">=3.10,<3.11"
-content-hash = "683be637a6a8a8c72ad256b5009268724ef81af2d7f53ee6c1cd9b41863f1272"
+content-hash = "4aa8543384cc30270d68de881c7a1ba0efa21d4f61406f29003c2de08b4e0f67"
diff --git a/pyproject.toml b/pyproject.toml
index 72ada4d..a9c6707 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -69,6 +69,7 @@ nflows = "^0.14"
 tensorboard = "^2.17.1"
 torch-geometric = "^2.5.3"
 lightning = "^2.4.0"
+gpytorch = "^1.13"
 
 
 [tool.poetry.group.dev.dependencies]