Remove deprecated torch._six imports

deepspeed/runtime/utils.py

@@ -16,7 +16,7 @@
 from bisect import bisect_left
 
 import torch
-from torch._six import inf
+from torch import inf
 from deepspeed import comm as dist
 
 from deepspeed.utils import groups, logger

deepspeed/runtime/zero/stage3.py

@@ -7,7 +7,7 @@
 import gc
 import collections
 from typing import Deque, Dict, Tuple
-from torch._six import inf
+from torch import inf
 
 from deepspeed.runtime import ZeROOptimizer
 from deepspeed.utils import logger

deepspeed/runtime/zero/stage_1_and_2.py

@@ -5,7 +5,7 @@
 import torch
 import os
 from deepspeed import comm as dist
-from torch._six import inf
+from torch import inf
 from packaging import version as pkg_version
 from collections import OrderedDict