From e961109f07a770a1587b7b6673bb0261cc56473d Mon Sep 17 00:00:00 2001 From: Lingxiao Zhang <823951506@qq.com> Date: Sat, 2 Jul 2022 16:54:38 +0800 Subject: [PATCH] Update train_encoder.py --- train_encoder.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/train_encoder.py b/train_encoder.py index 67fea3f..f5301bf 100644 --- a/train_encoder.py +++ b/train_encoder.py @@ -59,7 +59,8 @@ def embed_one_example(args, path, stylegan_encoder, g_all, upsamplers, label_im_tensor = label_im_tensor * 2.0 - 1.0 label_im_tensor = label_im_tensor.unsqueeze(0) latent_in = stylegan_encoder(label_im_tensor) - im_out_wo_encoder, _ = latent_to_image(g_all, upsamplers, latent_in, + with torch.no_grad(): + im_out_wo_encoder, _ = latent_to_image(g_all, upsamplers, latent_in, process_out=True, use_style_latents=True, return_only_im=True) @@ -173,7 +174,7 @@ def main(args, resume): ).to(device) stylegan_encoder = FPNEncoder(3, n_latent=args['n_latent'], only_last_layer=args['use_w'], same_view_code=args['same_view_code']) stylegan_encoder = stylegan_encoder.to(device) - + stylegan_encoder.eval() if resume != "": stylegan_encoder.load_state_dict(torch.load(resume, map_location=device)['model_state_dict'])