How to write an optimiser for StyleGAN2 interpolation?

I would like to interpolate two images using StyleGAN2-ADA-PyTorch from NVLabs. For the sake of simplicity, it can be said that with two images of different persons I want to create a third image depicting a third person, with a body from the first image, and their head from the second. I also have corresponding w-vectors for the two images ready at hand.

# G is a generative model in line with StyleGAN2, trained to output 512x512 images.
# Latents shape is [1, 16, 512]
G = G.eval().requires_grad_(False).to(device) # type: ignore
num_ws = G.mapping.num_ws # 16
w_dim = G.mapping.w_dim # 512

# Segmentation network is used to extract important parts from images
segmentation_dnn = segmentation_dnn.to(device)

# Source images are represented as latent vectors. I use G to generate actual images:
image_body = image_from_output(G.synthesis(w_body, noise_mode='const'))
image_head = image_from_output(G.synthesis(w_head, noise_mode='const'))

# Custom function is applied to source images, creating masked images. 
# In masked images, only head or body is present (and the rest is filled with white pixels)
image_body_masked = apply_segmentation_mask(image_body, segmentation_dnn, select='body')
image_head_masked = apply_segmentation_mask(image_head, segmentation_dnn, select='head')

In order to compare similarity of any two images, I use VGGLos

# VGG16 is used as a feature extractor to evaluate image similarity
url = 'https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/metrics/vgg16.pt'
with dnnlib.util.open_url(url) as f:
    vgg16 = torch.jit.load(f).eval().to(device)

class VGGLoss(nn.Module):
    def __init__(self, device, vgg):
        super().__init__()
        
        for param in self.parameters():
            param.requires_grad = False

        self.criterion = nn.L1Loss().to(device)
        
    def forward(self, source, target):
        loss = 0 
        source_features = self.vgg(source, resize_images=False, return_lpips=True)
        target_features = self.vgg(target, resize_images=False, return_lpips=True)
        loss += self.criterion(source, target)
            
        return loss 

vgg_loss = VGGLoss(device, vgg=vgg16)

Now, I want to interpolate image_body and image_head, creating image_target. To do this, I need to find latent representation of image_target in the latent space of StyleGAN2 Crudely, we can use optimize for a coefficient query_opt to partially include latents from image_body and image_head: w_target = w_body + (query_opt * (w_head - w_person))

query_opt = torch.randn([1, num_ws, 1], dtype=torch.float32, device=device, requires_grad=True)
optimizer = torch.optim.Adam(query_opt, betas=(0.9, 0.999), lr=initial_learning_rate)

w_out = []
for step in num_steps:
    # Learning rate schedule.
    t = step / num_steps
    lr_ramp = min(1.0, (1.0 - t) / lr_rampdown_length)
    lr_ramp = 0.5 - 0.5 * np.cos(lr_ramp * np.pi)
    lr_ramp = lr_ramp * min(1.0, t / lr_rampup_length)
    lr = initial_learning_rate * lr_ramp
    for param_group in optimizer.param_groups:
        param_group['lr'] = lr

    # Synth image from w_target using query_opt.
    # This interpolation formula is an important step, and I think my math might be out of order up here

    w_target = w_body + (query_opt * (w_head - w_person))
    image_target = image_from_output(G.synthesis(ws, noise_mode='const'))
    image_target_body_masked = apply_segmentation_mask(image_target, segmentation_dnn, select='body')
    image_target_head_masked = apply_segmentation_mask(image_target, segmentation_dnn, select='head')
    loss = vgg_loss(image_body_masked, image_target_body_masked) + vgg_loss(image_head_masked, image_target_head_masked)
    
    # Step
    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    optimizer.step()
    logprint(f'step {step+1:>4d}/{num_steps}: loss {float(loss):<5.2f}')

    # Save current w_target
    w_out[step] = w_target.detach()

I can't figure out how to make my optimizer actually target query_opt in such a way that combined VGGLoss is actually optimized for. I must be missing something in my PyTorch code, or maybe even in the main interpolation formula.