KeyError: 'G'

[Vadim2S]

I am do some experiments with StyleGAN_Human model in GUI version. Visualizer_drag.py much faster than web on my machine. And i am get error:

Traceback (most recent call last):
  File "e:\_Vadim_Work\2023_2\DragGAN\viz\renderer.py", line 142, in get_network
    data = legacy.load_network_pkl(f)
  File "e:\_Vadim_Work\2023_2\DragGAN\legacy.py", line 43, in load_network_pkl
    assert isinstance(data['G'], torch.nn.Module)
KeyError: 'G'

After research i am found cause: models like FFHQ contains G, D, G_ema generators, but StyleGAN_Human contains only G_ema. And therefore assertion in legacy.py rised

    # Validate contents.
    assert isinstance(data['G'], torch.nn.Module)
    assert isinstance(data['D'], torch.nn.Module)
    assert isinstance(data['G_ema'], torch.nn.Module)

However at this point we need only G_ema generator not others! Suggestion: 1) Change viz/renderer.py as:

def get_network(self, pkl, key, **tweak_kwargs):
...
    data = legacy.load_network_pkl(f, key)

2) Change legacy.py as:

def load_network_pkl(f, key=None, force_fp16=False):
...
    # Validate contents.
    if key is None or key == 'G':
        assert isinstance(data['G'], torch.nn.Module)
    if key is None or key == 'D':
        assert isinstance(data['D'], torch.nn.Module)
    if key is None or key == 'G_ema':
        assert isinstance(data['G_ema'], torch.nn.Module)

or something like this.