Fix UT (#449)
Fix UT (#449)

Co-authored-by: duanbing.0 <2bf81832effbd22c255597d48fa1abe8aa91bd79@bytedance.com>
bytedance/fedlearner,bytedance/fedlearner,bytedance/fedlearner,bytedance/fedlearner,bytedance/fedlearner,bytedance/fedlearner,bytedance/fedlearner