added torch to ml_defaults

CrayLabs · Aug 20, 2024 · e7c1882 · e7c1882
1 parent 9b6681c
commit e7c1882
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/smartsim/_core/_install/buildenv.py b/smartsim/_core/_install/buildenv.py
@@ -211,6 +211,7 @@ def ml_extras_required(self) -> t.Dict[str, t.List[str]]:
         """Optional ML/DL dependencies we suggest for the user.
         """
         ml_defaults = {
+            "torch": self.TORCH,
             "tensorflow": self.TENSORFLOW,
             "onnx": self.ONNX,
             "skl2onnx": "1.16.0",