add sklearn dependency

mir-group · Oct 3, 2021 · 420db74 · 420db74
1 parent 889f74c
commit 420db74
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/setup.py b/setup.py
@@ -29,13 +29,14 @@
         "numpy",
         "ase",
         "tqdm",
-        "torch>=1.8",
+        "torch>=1.8",  # torch.fx added in 1.8
         "torch_scatter",
         "e3nn>=0.3.3",
         "pyyaml",
         "contextlib2;python_version<'3.7'",  # backport of nullcontext
-        "typing_extensions;python_version<'3.8'",
+        "typing_extensions;python_version<'3.8'",  # backport of Final
         "torch-runstats>=0.2.0",
+        "scikit_learn",  # for GaussianProcess for per-species statistics
     ],
     zip_safe=True,
 )