diff --git a/pyproject.toml b/pyproject.toml
index 1fa82d12..a9356ffd 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -58,8 +58,10 @@ docs = [
     "autodoc_pydantic==2.2.0",
     "furo==2024.8.6",
     "ipython==8.34.0",
+    "ipykernel==6.30.1",
     "jsonschema[format]",
     "jupyterlab==4.3.4",
+    "jupyter-core==5.8.1",
     "jupytext==1.16.7",
     "myst_parser==4.0.0",
     "nbsphinx>=0.9.7",
diff --git a/tests/models/test_fairchem.py b/tests/models/test_fairchem.py
index bac01aa0..f445a987 100644
--- a/tests/models/test_fairchem.py
+++ b/tests/models/test_fairchem.py
@@ -219,16 +219,16 @@ def test_empty_batch_error() -> None:
 @pytest.mark.skipif(
     get_token() is None, reason="Requires HuggingFace authentication for UMA model access"
 )
-def test_load_from_checkpoint_path(device: torch.device, dtype: torch.dtype) -> None:
+def test_load_from_checkpoint_path() -> None:
     """Test loading model from a saved checkpoint file path."""
     checkpoint_path = pretrained_checkpoint_path_from_name("uma-s-1")
     loaded_model = FairChemModel(
-        model=str(checkpoint_path), task_name="omat", cpu=device.type == "cpu"
+        model=str(checkpoint_path), task_name="omat", cpu=DEVICE == "cpu"
     )
 
     # Verify the loaded model works
     system = bulk("Si", "diamond", a=5.43)
-    state = ts.io.atoms_to_state([system], device=device, dtype=dtype)
+    state = ts.io.atoms_to_state([system], device=DEVICE, dtype=DTYPE)
     results = loaded_model(state)
 
     assert "energy" in results