colossalai.utils.cuda

colossalai.utils.cuda.set_to_cuda(models)

Send model to gpu.

Parameters

models – nn.module or a list of module

colossalai.utils.cuda.get_current_device()

Returns the index of a currently selected device (gpu/cpu).

colossalai.utils.cuda.synchronize()

Similar to cuda.synchronize(). Waits for all kernels in all streams on a CUDA device to complete.

colossalai.utils.cuda.empty_cache()

Similar to cuda.empty_cache() Releases all unoccupied cached memory currently held by the caching allocator.