wrong judgment condition of train epochs

[j-cyoung]

As writen in the Branch_CIFAR_10/squeeze/squeeze.sh and the paper supplemental, the total epochs of pretrained model is 5

python train.py --model ResNet18 --dataset CIFAR-10 --data_path /path/to/cifar-10 --train_epochs 5

python train.py --model MobileNetV2 --dataset CIFAR-10 --data_path /path/to/cifar-10 --train_epochs 5

python train.py --model ShuffleNetV2_0_5 --dataset CIFAR-10 --data_path /path/to/cifar-10 --train_epochs 5

python train.py --model WRN_16_2 --dataset CIFAR-10 --data_path /path/to/cifar-10 --train_epochs 5

python train.py --model ConvNetW128 --dataset CIFAR-10 --data_path /path/to/cifar-10 --train_epochs 5

However it seems that the judgment condition in Branch_CIFAR_10/squeeze/train.py (line 121) cannot save the corresponding epoch:

        if e % 10 == 0:
            state_dict = model.state_dict()
            torch.save(state_dict,os.path.join(save_dir,f"squeeze_{args.model}.pth"))