oneflow.cuda¶
Returns a bool indicating if CUDA is currently available. |
|
Returns the number of GPUs available. |
|
Returns local rank as device index. |
|
Sets the current device. |
|
Waits for all kernels in all streams on a CUDA device to complete. |
|
Gets the properties of a device. |
|
Gets the cuda capability of a device. |
|
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¶
Sets the seed for generating random numbers on all GPUs. |
|
Sets the seed for generating random numbers for the current GPU. |
|
Returns the random number generator state of the specified GPU as a ByteTensor. |
|
Returns a list of ByteTensor representing the random number states of all devices. |
GPU tensor¶
The tensortype oneflow.cuda.HalfTensor is not available. |
|
The tensortype oneflow.cuda.FloatTensor is not available. |
|
The tensortype oneflow.cuda.DoubleTensor is not available. |
|
The tensortype oneflow.cuda.BoolTensor is not available. |
|
The tensortype oneflow.cuda.ByteTensor is not available. |
|
The tensortype oneflow.cuda.CharTensor is not available. |
|
The tensortype oneflow.cuda.IntTensor is not available. |
|
The tensortype oneflow.cuda.LongTensor is not available. |
Memory management¶
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. |