hiro/ray
1
0
Fork 0
mirror of https://github.com/vale981/ray synced 2025-03-06 10:31:39 -05:00
Code Issues Projects Releases Packages Wiki Activity Actions

[Train] Decorate get_device with PublicAPI (#22024)

Browse source 
* Decorate `get_device` with `PublicAPI`

* Add documentation

* Update api.rst
This commit is contained in:
Balaji Veeramani 2022-02-01 08:18:47 -08:00 • committed by GitHub
parent b51b5afaea
commit 7dcb0b6af6
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 1 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
python/ray/train/torch.py
View file

@ -210,6 +210,7 @@ class _WrappedDataLoader(DataLoader):
yield self._move_to_device(item)
@PublicAPI(stability="beta")
def get_device() -> torch.device:
"""Gets the correct torch device to use for training."""
if torch.cuda.is_available():