flytekit.extras.accelerators.GPUAccelerator

class flytekit.extras.accelerators.GPUAccelerator(device)

Class that represents a GPU accelerator. The class can be instantiated with any valid GPU device name, but it is recommended to use one of the pre-defined constants below, as name has to match the name of the device configured on the cluster.

Methods

Parameters:

device (str)

to_flyte_idl()

Return type:

GPUAccelerator