forked from mindspore-Ecosystem/mindspore
!7149 add exception for mixed precision when device is CPU
Merge pull request !7149 from wukesong/lenet-mix-fix
This commit is contained in:
commit
d1e029013e
|
@ -60,7 +60,11 @@ if __name__ == "__main__":
|
|||
config_ck = CheckpointConfig(save_checkpoint_steps=cfg.save_checkpoint_steps,
|
||||
keep_checkpoint_max=cfg.keep_checkpoint_max)
|
||||
ckpoint_cb = ModelCheckpoint(prefix="checkpoint_lenet", directory=args.ckpt_path, config=config_ck)
|
||||
model = Model(network, net_loss, net_opt, metrics={"Accuracy": Accuracy()}, amp_level="O2")
|
||||
|
||||
if args.device_target == "CPU":
|
||||
model = Model(network, net_loss, net_opt, metrics={"Accuracy": Accuracy()})
|
||||
else:
|
||||
model = Model(network, net_loss, net_opt, metrics={"Accuracy": Accuracy()}, amp_level="O2")
|
||||
|
||||
print("============== Starting Training ==============")
|
||||
model.train(cfg['epoch_size'], ds_train, callbacks=[time_cb, ckpoint_cb, LossMonitor()],
|
||||
|
|
Loading…
Reference in New Issue