From aa0b7c70450a5ff65956e4964f48f33649a25233 Mon Sep 17 00:00:00 2001 From: Eduardo Apolinario Date: Thu, 25 Jul 2024 16:13:56 -0400 Subject: [PATCH] Add nvidia-l4 gpu accelerator Signed-off-by: Eduardo Apolinario --- 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 `_