colossalai.utils.cuda

colossalai.utils.cuda.set_to_cuda(models)[source]

Send model to gpu.

Parameters: models – nn.module or a list of module

colossalai.utils.cuda.get_current_device()[source]: Returns the index of a currently selected device (gpu/cpu).

colossalai.utils.cuda.synchronize()[source]: Similar to cuda.synchronize(). Waits for all kernels in all streams on a CUDA device to complete.

colossalai.utils.cuda.empty_cache()[source]: Similar to cuda.empty_cache() Releases all unoccupied cached memory currently held by the caching allocator.