Fix AssertionError

#25
NVlabs · Nov 3, 2020 · 1d29bb0 · 1d29bb0
1 parent 8c3fd8b
commit 1d29bb0
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/generate.py b/generate.py
@@ -33,7 +33,7 @@ def generate_images(network_pkl, seeds, truncation_psi, outdir, class_idx, dlate
     if dlatents_npz is not None:
         print(f'Generating images from dlatents file "{dlatents_npz}"')
         dlatents = np.load(dlatents_npz)['dlatents']
-        assert dlatents.shape[1:] == (18, 512) # [N, 18, 512]
+        assert dlatents.shape[-1] == 512 # [N, M, 512]
         imgs = Gs.components.synthesis.run(dlatents, output_transform=dict(func=tflib.convert_images_to_uint8, nchw_to_nhwc=True))
         for i, img in enumerate(imgs):
             fname = f'{outdir}/dlatent{i:02d}.png'