From 27e160c63e09501b4289f1719120ab1f12dded74 Mon Sep 17 00:00:00 2001
From: Julia Putko <juliaputko@gmail.com>
Date: Tue, 20 Aug 2024 16:08:44 -0500
Subject: [PATCH] added cpu suffix and cuda suffix for torch defaults

---
 smartsim/_core/_install/buildenv.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/smartsim/_core/_install/buildenv.py b/smartsim/_core/_install/buildenv.py
index 4212a223a..72e09a4b9 100644
--- 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