pytorch/docs/source/mtia.rst
Trung Truong cc365fdd7b [MTIA] Support torch.cuda.get_device_capability equivalent API on MTIA (#135889)
Summary:
Mirror `get_device_capability` on MTIA per https://fburl.com/gdoc/p4lo5avn

At the moment, both the major and minor version are just 0

Test Plan:
Unit test: `buck2 test //mtia/host_runtime/torch_mtia/tests:test_torch_mtia_api`

https://www.internalfb.com/intern/testinfra/testconsole/testrun/1688850109958190/

Differential Revision: D62595296

Pull Request resolved: https://github.com/pytorch/pytorch/pull/135889
Approved by: https://github.com/egienvalue
2024-09-17 17:42:56 +00:00

40 lines
685 B
ReStructuredText

torch.mtia
===================================
The MTIA backend is implemented out of the tree, only interfaces are be defined here.
.. automodule:: torch.mtia
.. currentmodule:: torch.mtia
.. autosummary::
:toctree: generated
:nosignatures:
StreamContext
current_device
current_stream
default_stream
device_count
init
is_available
is_initialized
memory_stats
get_device_capability
set_device
set_stream
stream
synchronize
device
set_rng_state
get_rng_state
DeferredMtiaCallError
Streams and events
------------------
.. autosummary::
:toctree: generated
:nosignatures:
Event
Stream