added cpu suffix and cuda suffix for torch defaults

CrayLabs · Aug 20, 2024 · 27e160c · 27e160c
1 parent ae32e53
commit 27e160c
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/smartsim/_core/_install/buildenv.py b/smartsim/_core/_install/buildenv.py
@@ -218,6 +218,8 @@ def ml_extras_required(self) -> t.Dict[str, t.List[str]]:
             "onnxmltools": "1.12.0",
             "scikit-learn": "1.3.2",
             "torchvision": "0.15.2",
+            "torch_cpu_suffix": "+cpu",
+            "torch_cuda_suffix": "+cu117",
         }
 
         # remove torch-related fields as they are subject to change