From e71a76009c54276eb7c7906c046122a4e236a31e Mon Sep 17 00:00:00 2001 From: MrGranddy Date: Sun, 8 Mar 2020 10:50:47 +0300 Subject: [PATCH] A fix to use 18x512 projected latent spaces --- model.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/model.py b/model.py index 4493cb6d..fb967622 100755 --- a/model.py +++ b/model.py @@ -516,8 +516,14 @@ def forward( if inject_index is None: inject_index = random.randint(1, self.n_latent - 1) - latent = styles[0].unsqueeze(1).repeat(1, inject_index, 1) - latent2 = styles[1].unsqueeze(1).repeat(1, self.n_latent - inject_index, 1) + if styles[0].ndim < 3: + latent = styles[0].unsqueeze(1).repeat(1, inject_index, 1) + else: + latent = styles[0][:,:inject_index, :] + if styles[1].ndim < 3: + latent2 = styles[1].unsqueeze(1).repeat(1, self.n_latent - inject_index, 1) + else: + latent2 = styles[1][:,inject_index:,:] latent = torch.cat([latent, latent2], 1)