mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2025-04-28 11:45:23 +00:00
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
This commit is contained in:
parent
a88bc828d5
commit
01066152f1
@ -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
|
||||
|
@ -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
|
||||
|
Loading…
Reference in New Issue
Block a user