parser = argparse.ArgumentParser()
parser.add_argument("--interrupt", action="store_true")
args = parser.parse_args()

callbacks = [keras.callbacks.BackupAndRestore(backup_dir=BACKUP_DIR)]
if args.interrupt:
    callbacks.append(InterruptAtEpoch(epoch_to_stop=3))

model = build_model()