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.

set_rng_state

Sets the random number generator state of the specified GPU.

set_rng_state_all

Sets the random number generator state 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.