4a5eef64d2a0ab5e1737ffbb345d9e72d2f1adc9,train.py,,,#,56

Before Change

vgg = nn.Sequential(*list(vgg.children())[:31])
network = net.Net(vgg, decoder)

for name in ["enc_1", "enc_2", "enc_3", "enc_4"]:
    for param in getattr(network, name).parameters():
        param.requires_grad = False
network.train()
network.cuda()

content_tf = train_transform()

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()

Instances