cuda.rst 1 KB
Newer Older
yuguo's avatar
yuguo committed
1
2
oneflow.cuda
===================================
yuguo's avatar
yuguo committed
3
4
5

.. The documentation is referenced from: https://pytorch.org/docs/1.10/cuda.html.

yuguo's avatar
yuguo committed
6
.. currentmodule:: oneflow.cuda
yuguo's avatar
yuguo committed
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57

.. autosummary::
    :toctree: generated
    :nosignatures:

    is_available
    device_count
    current_device
    set_device
    synchronize
    get_device_properties
    get_device_capability
    get_device_name

.. note::
   The :attr:`current_device` returns local rank as device index. It is different from the 'torch.current_device()' in PyTorch.


Random Number Generator
-------------------------
.. autosummary::
    :toctree: generated
    :nosignatures:

    manual_seed_all
    manual_seed


GPU tensor
-----------------------------
.. autosummary::
    :toctree: generated
    :nosignatures:

    HalfTensor
    FloatTensor
    DoubleTensor
    BoolTensor
    ByteTensor
    CharTensor
    IntTensor
    LongTensor

Memory management
-----------------------------
.. autosummary::
    :toctree: generated
    :nosignatures:

    empty_cache