From d7fadc4731c41f675007112a801979e5d507222a Mon Sep 17 00:00:00 2001
From: flybird11111 <1829166702@qq.com>
Date: Mon, 20 May 2024 06:38:41 +0000
Subject: [PATCH] fix

---
 colossalai/checkpoint_io/utils.py | 5 -----
 1 file changed, 5 deletions(-)

diff --git a/colossalai/checkpoint_io/utils.py b/colossalai/checkpoint_io/utils.py
index 6197be9d1c8d..7acada0c55fa 100644
--- a/colossalai/checkpoint_io/utils.py
+++ b/colossalai/checkpoint_io/utils.py
@@ -111,11 +111,6 @@ def search_tp_partition_dim(current_shape: torch.Size, original_shape: torch.Siz
         if length > current_shape[dim]:
             partition_dim = dim
             break
-    if partition_dim is not None:
-        assert (
-            original_shape[partition_dim] == tp_size * current_shape[partition_dim]
-        ), f"The parameter isn't evenly distributed among tensor parallel group: \
-                shape before sharding {original_shape}, shape after sharding {current_shape}"
 
     return partition_dim