oneflow.cuda

is_available	Returns a bool indicating if CUDA is currently available.
device_count	Returns the number of GPUs available.
current_device	Returns local rank as device index.
set_device	Sets the current device.
synchronize	Waits for all kernels in all streams on a CUDA device to complete.
get_device_properties	Gets the properties of a device.
get_device_capability	Gets the cuda capability of a device.
get_device_name	Gets the name of a device.

Note

The current_device returns local rank as device index. It is different from the ‘torch.current_device()’ in PyTorch.

Random Number Generator

manual_seed_all	Sets the seed for generating random numbers on all GPUs.
manual_seed	Sets the seed for generating random numbers for the current GPU.
get_rng_state	Returns the random number generator state of the specified GPU as a ByteTensor.
get_rng_state_all	Returns a list of ByteTensor representing the random number states of all devices.

GPU tensor

HalfTensor	The tensortype oneflow.cuda.HalfTensor is not available.
FloatTensor	The tensortype oneflow.cuda.FloatTensor is not available.
DoubleTensor	The tensortype oneflow.cuda.DoubleTensor is not available.
BoolTensor	The tensortype oneflow.cuda.BoolTensor is not available.
ByteTensor	The tensortype oneflow.cuda.ByteTensor is not available.
CharTensor	The tensortype oneflow.cuda.CharTensor is not available.
IntTensor	The tensortype oneflow.cuda.IntTensor is not available.
LongTensor	The tensortype oneflow.cuda.LongTensor is not available.

Memory management

empty_cache

Releases all unoccupied cached memory currently held by the caching allocators of all OneFlow streams so those can be re-allocated in OneFlow streams or other GPU application and visible in nvidia-smi.