From 537a3d719af4c732276335009a4f3f16e819212e Mon Sep 17 00:00:00 2001 From: Eduardo Apolinario <653394+eapolinario@users.noreply.github.com> Date: Fri, 26 Jul 2024 11:22:29 -0400 Subject: [PATCH] Add nvidia-l4 gpu accelerator (#2608) Signed-off-by: Eduardo Apolinario Co-authored-by: Eduardo Apolinario Signed-off-by: mao3267 --- flytekit/extras/accelerators.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/flytekit/extras/accelerators.py b/flytekit/extras/accelerators.py index 8a9d3e56a5..7cc3bb6bd5 100644 --- a/flytekit/extras/accelerators.py +++ b/flytekit/extras/accelerators.py @@ -133,7 +133,11 @@ def to_flyte_idl(self) -> tasks_pb2.GPUAccelerator: #: use this constant to specify that the task should run on an #: `NVIDIA L4 Tensor Core GPU `_ -L4 = GPUAccelerator("nvidia-l4-vws") +L4 = GPUAccelerator("nvidia-l4") + +#: use this constant to specify that the task should run on an +#: `NVIDIA L4 Tensor Core GPU `_ +L4_VWS = GPUAccelerator("nvidia-l4-vws") #: use this constant to specify that the task should run on an #: `NVIDIA Tesla K80 GPU `_