wehub-resource-sync
55 lines
1.3 KiB
Python
55 lines
1.3 KiB
Python
# isort: off
|
|
try:
|
|
import torch # noqa: F401
|
|
except ModuleNotFoundError:
|
|
raise ModuleNotFoundError(
|
|
"PyTorch isn't installed. To install PyTorch, run 'pip install torch'"
|
|
)
|
|
# isort: on
|
|
|
|
from ray.train.torch.config import TorchConfig
|
|
from ray.train.torch.torch_checkpoint import TorchCheckpoint
|
|
from ray.train.torch.torch_trainer import TorchTrainer
|
|
from ray.train.torch.train_loop_utils import (
|
|
accelerate,
|
|
backward,
|
|
enable_reproducibility,
|
|
get_device,
|
|
get_devices,
|
|
prepare_data_loader,
|
|
prepare_model,
|
|
prepare_optimizer,
|
|
)
|
|
from ray.train.v2._internal.constants import is_v2_enabled
|
|
|
|
if is_v2_enabled():
|
|
from ray.train.v2.torch.torch_trainer import TorchTrainer # noqa: F811
|
|
from ray.train.v2.torch.train_loop_utils import ( # noqa: F811
|
|
accelerate,
|
|
backward,
|
|
enable_reproducibility,
|
|
get_device,
|
|
get_devices,
|
|
prepare_data_loader,
|
|
prepare_model,
|
|
prepare_optimizer,
|
|
)
|
|
|
|
|
|
__all__ = [
|
|
"TorchTrainer",
|
|
"TorchCheckpoint",
|
|
"TorchConfig",
|
|
"accelerate",
|
|
"get_device",
|
|
"get_devices",
|
|
"prepare_model",
|
|
"prepare_optimizer",
|
|
"prepare_data_loader",
|
|
"backward",
|
|
"enable_reproducibility",
|
|
]
|
|
|
|
|
|
# DO NOT ADD ANYTHING AFTER THIS LINE.
|