1
0

__init__.py 604 B

1234567891011121314151617181920212223
  1. from typing import Optional
  2. import torch
  3. from aphrodite.common.utils import is_tpu
  4. from .interface import Platform, PlatformEnum, UnspecifiedPlatform
  5. current_platform: Optional[Platform]
  6. if torch.version.cuda is not None:
  7. from .cuda import CudaPlatform
  8. current_platform = CudaPlatform()
  9. elif torch.version.hip is not None:
  10. from .rocm import RocmPlatform
  11. current_platform = RocmPlatform()
  12. elif is_tpu():
  13. from .tpu import TpuPlatform
  14. current_platform = TpuPlatform()
  15. else:
  16. current_platform = UnspecifiedPlatform()
  17. __all__ = ['Platform', 'PlatformEnum', 'current_platform']