diff --git a/model_zoo/official/cv/mobilenetv1/train.py b/model_zoo/official/cv/mobilenetv1/train.py index aa29eedc464..8da7bed47a2 100755 --- a/model_zoo/official/cv/mobilenetv1/train.py +++ b/model_zoo/official/cv/mobilenetv1/train.py @@ -143,7 +143,7 @@ if __name__ == '__main__': time_cb = TimeMonitor(data_size=step_size) loss_cb = LossMonitor() cb = [time_cb, loss_cb] - if config.save_checkpoint: + if config.save_checkpoint and device_id % min(8, args_opt.device_num) == 0: config_ck = CheckpointConfig(save_checkpoint_steps=config.save_checkpoint_epochs * step_size, keep_checkpoint_max=config.keep_checkpoint_max) ckpt_cb = ModelCheckpoint(prefix="mobilenetv1", directory=ckpt_save_dir, config=config_ck)