I what to run this in heterogeneous data, but there are some errors

[LittWhite-WANG]

When i wangt to use the model of simple-cnn under heterogeneous data, there are some errors , my commond is :
python dist_skew_main.py --model=simple-cnn --dataset=cifar10 --lr=0.01 --retrain_lr=0.01 --batch-size=64 --epochs=10 --retrain_epochs=10 --n_nets=10 --partition=hetero-dir --comm_type=fedma --comm_round=10 --retrain=True --rematching=True

and the error is :
Traceback (most recent call last):
File "dist_skew_main.py", line 1181, in
args.partition, args.n_nets, args_alpha, args=args)
File "/home/wjj/three/FedMA-master/utils.py", line 274, in partition_data_dist_skew
traindata_cls_counts = record_net_data_stats(y_train, net_dataidx_map, logdir)
UnboundLocalError: local variable 'net_dataidx_map' referenced before assignment

the code is

Can you tell me how to solve this problem