parser = argparse.ArgumentParser()
// Basic options
parser.add_argument("--gpu", type=int, default=-1)
parser.add_argument("--content_dir", type=str, required=True,
help="Directory path to a batch of content images")
parser.add_argument("--style_dir", type=str, required=True,
help="Directory path to a batch of style images")
After Change
if (i + 1) % args.save_model_interval == 0 or (i + 1) == args.max_iter:
state_dict = net.decoder.state_dict()
for key in state_dict.keys():
state_dict[key] = state_dict[key].to(torch.device("cpu"))
torch.save(state_dict,
"{:s}/decoder_iter_{:d}.pth.tar".format(args.save_dir,
i + 1))
writer.close()
