Restrict n_samples of the inference script according to number of lat…

…ents.
omertov · May 19, 2022 · 6525c5f · 6525c5f
1 parent 6f8d68a
commit 6525c5f
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/scripts/inference.py b/scripts/inference.py
@@ -101,7 +101,7 @@ def generate_inversions(args, g, latent_codes, is_cars):
     print('Saving inversion images')
     inversions_directory_path = os.path.join(args.save_dir, 'inversions')
     os.makedirs(inversions_directory_path, exist_ok=True)
-    for i in range(args.n_sample):
+    for i in range(min(args.n_sample, len(latents))):
         imgs, _ = g([latent_codes[i].unsqueeze(0)], input_is_latent=True, randomize_noise=False, return_latents=True)
         if is_cars:
             imgs = imgs[:, :, 64:448, :]