style: polish code

Author: cwher

colossalai/shardformer/policies/base_policy.py

@@ -242,4 +242,4 @@ def get_stage_index(
             end_idx = num_layers_per_stage_accumulated[stage + model_chunk * num_stages + 1]
             stage_indices.append([start_idx, end_idx])
 
-        return stage_indices[0] if num_model_chunks == 1 else stage_indices
+        return stage_indices[0] if num_model_chunks == 1 else stage_indices

tests/test_shardformer/test_layer/test_dist_crossentropy.py

@@ -38,9 +38,10 @@ def check_dist_crossentropy(rank, world_size, port, ignore_index):
         org_loss, dist_loss, atol=1e-5
     ), f"dist cross entropy loss is not equal to orgin loss\n{org_loss}\n{dist_loss}"
 
-
     target_grad = torch.chunk(pred.grad, world_size, dim=-1)[rank]
-    assert torch.allclose(target_grad, dist_pred.grad), f"dist grad is not equal to orgin grad\n{target_grad}\n{dist_pred.grad}"
+    assert torch.allclose(
+        target_grad, dist_pred.grad
+    ), f"dist grad is not equal to orgin grad\n{target_grad}\n{dist_pred.grad}"
 
 
 @pytest.mark.dist

tests/test_shardformer/test_model/test_shard_gptj.py

@@ -207,6 +207,7 @@ def check_gptj_3d(rank, world_size, port):
     colossalai.launch(config={}, rank=rank, world_size=world_size, host="localhost", port=port, backend="nccl")
     run_gptj_3d_test()
 
+
 @pytest.mark.skip("TODO check_gptj has something wrong.")
 @pytest.mark.dist
 @rerun_if_address_is_in_use()