Issue in pruning a .pt StyleGAN2-ADA-PyTorch model

[imjohnzakkam]

Hi, I have been trying to prune a StyleGAN2-ADA-PyTorch model (which is saved in .pt format) which by default saves in .pkl format but, I have changed it to .pt (by using torch.save() and saving needed dicts) as the implementation needs .pt format, now the issue is with the file prune.py which is giving me an error as mentioned below:

Traceback (most recent call last):
  File "prune.py", line 36, in <module>
    g_ema = Build_Generator_From_Dict(model_dict['g_ema'], size=args.generated_img_size).to(device)
  File "/home/notebook/code/personal/605241/content-aware-gan-compression/Util/network_util.py", line 101, in Build_Generator_From_Dict
    net_shape = Get_Network_Shape(model_dict)
  File "/home/notebook/code/personal/605241/content-aware-gan-compression/Util/network_util.py", line 37, in Get_Network_Shape
    num_channels = [model_dict[key].shape[2] for key in conv_key_list] # from start to end
  File "/home/notebook/code/personal/605241/content-aware-gan-compression/Util/network_util.py", line 37, in <listcomp>
    num_channels = [model_dict[key].shape[2] for key in conv_key_list] # from start to end
KeyError: 'conv1.conv.weight'

The referenced error redirects to /Util/network_util.py 's function Get_Conv_Kernel_Key() which has CONV1_KEY already set to 'conv1.conv.weight', but my model's generator dosen't even have a key as this. I am not sure how to make this work for my model, can you please check this out?