utils.py 409 B

12345678910111213141516
  1. '''Copyright The Microsoft DeepSpeed Team'''
  2. import torch
  3. def is_torch_elastic_compatible():
  4. '''
  5. Helper to lookup torch version. Elastic training is
  6. introduced in 1.11.x
  7. '''
  8. TORCH_MAJOR = int(torch.__version__.split('.')[0])
  9. TORCH_MINOR = int(torch.__version__.split('.')[1])
  10. if TORCH_MAJOR == 1 and TORCH_MINOR >= 11:
  11. return True
  12. else:
  13. return False