From 8e3f66a0d1ab44ecbee5917d15320e33ad118b5f Mon Sep 17 00:00:00 2001
From: Boyuan Yao <70263930+Cypher30@users.noreply.github.com>
Date: Fri, 17 Feb 2023 10:26:07 +0800
Subject: [PATCH] [zero] fix wrong import (#2777)

---
 colossalai/zero/sharded_optim/_utils.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/colossalai/zero/sharded_optim/_utils.py b/colossalai/zero/sharded_optim/_utils.py
index 68928b232..9ca2fdf5a 100644
--- a/colossalai/zero/sharded_optim/_utils.py
+++ b/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 import inf
 from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors
 
 from colossalai.tensor import ColoParameter