warp_drive.training.utils.single_device_context package

Submodules

warp_drive.training.utils.single_device_context.device_context_numba module

class warp_drive.training.utils.single_device_context.device_context_numba.NumbaSingleDeviceContext

Bases: object

one single GPU hardware context available for both Numba and PyTorch

clear_context()

property context

init_context(device_id=None)

warp_drive.training.utils.single_device_context.device_context_pycuda module

class warp_drive.training.utils.single_device_context.device_context_pycuda.PyCUDASingleDeviceContext

Bases: object

one single GPU hardware context available for both PyCuda and PyTorch

property context

init_context(device_id=None)

warp_drive.training.utils.single_device_context.device_context_pycuda.clear_context(context)

Module contents