Don't use `torch._six` (#2775)

* Don't use `torch._six`

This is a private API which is gone after https://github.com/pytorch/pytorch/pull/94709

* Update common.py

colossalai/utils/common.py

@@ -11,7 +11,7 @@ from typing import Callable, Dict, List, Optional, Union
 
 import torch
 import torch.distributed as dist
-from torch._six import inf
+from torch import inf
 from torch.nn.parameter import Parameter
 
 from colossalai.constants import IS_TENSOR_PARALLEL, NUM_PARTITIONS, TENSOR_PARALLEL_ATTRIBUTES

colossalai/zero/sharded_optim/_utils.py

@@ -3,7 +3,7 @@ from typing import Optional
 
 import torch
 import torch.distributed as dist
-from torch._six import inf
+from torch.six import inf
 from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors
 
 from colossalai.tensor import ColoParameter