Unable to run LSUN bedroom trained in Stylegan3

[AbhinavaShukla23]

Hi, I am trying to use pretrained pkl file fof LSUN bedroom trained in stylegan (https://drive.google.com/uc?id=1MOSKeGF0FJcivpBI7s63V9YHloUTORiF) refer to https://github.com/NVlabs/stylegan. However, when ran in Stylegan3, I am getting below error:

Loading networks from "/content/gdrive/MyDrive/diffusion-stylegan2-lsun-bedroom.pkl"... Traceback (most recent call last): File "/content/stylegan3/gen_images.py", line 143, in generate_images() # pylint: disable=no-value-for-parameter File "/usr/local/lib/python3.9/dist-packages/click/core.py", line 1130, in call return self.main(args, kwargs) File "/usr/local/lib/python3.9/dist-packages/click/core.py", line 1055, in main rv = self.invoke(ctx) File "/usr/local/lib/python3.9/dist-packages/click/core.py", line 1404, in invoke return ctx.invoke(self.callback, ctx.params) File "/usr/local/lib/python3.9/dist-packages/click/core.py", line 760, in invoke return __callback(args, **kwargs) File "/content/stylegan3/gen_images.py", line 108, in generate_images G = legacy.load_network_pkl(f)['G_ema'].to(device) # type: ignore File "/content/stylegan3/legacy.py", line 23, in load_network_pkl data = _LegacyUnpickler(f).load() File "/content/stylegan3/legacy.py", line 69, in find_class return super().find_class(module, name) AttributeError: Can't get attribute 'Identity' on <module 'training.augment' from '/content/stylegan3/training/augment.py'>

I am able to run the given network pkl files without any issue. I am running in google colab.