text: Title: Issues with Running stylegan2_pytorch in gpu settings on colab notebook
  Body: It keeps raising AttributeError whenever torch_stylegan2.load_network_pkl() is called.
  
  The code that I have run:
  
  `import stylegan2_pytorch as torch_stylegan2
  !mkdir networks
  !gdown https://nvlabs-fi-cdn.nvidia.com/stylegan2/networks/stylegan2-ffhq-config-f.pkl -O networks/stylegan2-ffhq-config-f.pkl
  network_pkl = "networks/stylegan2-ffhq-config-f.pkl"
  model_weights = torch_stylegan2.load_network_pkl(network_pkl)
  model = StyleGAN2(state_dict=model_weights)
  model.eval()`
  
  
  
  
  The error that reads,
  
  `mkdir: cannot create directory ‘networks’: File exists   Downloading... From: https://nvlabs-fi cdn.nvidia.com/stylegan2/networks/stylegan2-ffhq-config-f.pkl To: /content/networks/stylegan2-ffhq-config-f.pkl
  100% 382M/382M [00:02<00:00, 136MB/s] ---------------------------------------------------------------------------
  
  AttributeError                            Traceback (most recent call last)
  [<ipython-input-13-7df4095cee9d>](https://localhost:8080/#) in <cell line: 9>()
        7 # Load the pre-trained model weights
        8 network_pkl = "networks/stylegan2-ffhq-config-f.pkl"
  ----> 9 model_weights = torch_stylegan2.load_network_pkl(network_pkl)
       10 
       11 model = StyleGAN2(state_dict=model_weights)
  AttributeError: module 'stylegan2_pytorch' has no attribute 'load_network_pkl'
  
  
  
  
  
  I have already installed the package by running ' !pip install --upgrade stylegan2_pytorch' on the notebook with the output as described in truncated format:
  
  `Requirement already satisfied: websockets in /usr/local/lib/python3.10/dist-packages (from aim->stylegan2_pytorch) (12.0)
  Requirement already satisfied: boto3 in /usr/local/lib/python3.10/dist-packages (from aim->stylegan2_pytorch) (1.34.82)
  Requirement already satisfied: base58==2.0.1 in /usr/local/lib/python3.10/dist-packages (from aimrecords==0.0.7->aim->stylegan2_pytorch) (2.0.1)
  Requirement already satisfied: six in /usr/local/lib/python3.10/dist-packages (from fire->stylegan2_pytorch) (1.16.0)
  Requirement already satisfied: termcolor in /usr/local/lib/python3.10/dist-packages (from fire->stylegan2_pytorch) (2.4.0)
  Requirement already satisfied: decorator>=3.4.2 in /usr/local/lib/python3.10/dist-packages (from retry->stylegan2_pytorch) (4.4.2)
  Requirement already satisfied: py<2.0.0,>=1.4.26 in /usr/local/lib/python3.10/dist-packages (from retry->stylegan2_pytorch) (1.11.0)
  Requirement already satisfied: Mako in /usr/local/lib/python3.10/dist-packages (from alembic<2,>=1.5.0->aim->stylegan2_pytorch) (1.3.3)
  Requirement already satisfied: cffi>=1.12 in /usr/local/lib/python3.10/dist-packages (from cryptography>=3.0->aim->stylegan2_pytorch) (1.16.0)
  Requirement already satisfied: pydantic!=1.8,!=1.8.1,!=2.0.0,!=2.0.1,!=2.1.0,<3.0.0,>=1.7.4 in /usr/local/lib/python3.10/dist-packages (from fastapi<1,>=0.69.0->aim->stylegan2_pytorch) (2.6.4)
  Requirement already satisfied: starlette<0.38.0,>=0.37.2 in /usr/local/lib/python3.10/dist-packages (from fastapi<1,>=0.69.0->aim->stylegan2_pytorch) (0.37.2)
  Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.10/dist-packages (from jinja2->torch->stylegan2_pytorch) (2.1.5)
  Requirement already satisfied: greenlet!=0.4.17 in /usr/local/lib/python3.10/dist-packages (from SQLAlchemy>=1.4.1->aim->stylegan2_pytorch) (3.0.3)
  Requirement already satisfied: h11>=0.8 in /usr/local/lib/python3.10/dist-packages (from uvicorn<1,>=0.12.0->aim->stylegan2_pytorch) (0.14.0)
  Requirement already satisfied: botocore<1.35.0,>=1.34.82 in /usr/local/lib/python3.10/dist-packages (from boto3->aim->stylegan2_pytorch) (1.34.82)
  Requirement already satisfied: jmespath<2.0.0,>=0.7.1 in /usr/local/lib/python3.10/dist-packages (from boto3->aim->stylegan2_pytorch) (1.0.1)
  Requirement already satisfied: s3transfer<0.11.0,>=0.10.0 in /usr/local/lib/python3.10/dist-packages (from boto3->aim->stylegan2_pytorch) (0.10.1)
  Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.10/dist-packages (from requests->aim->stylegan2_pytorch) (3.3.2)
  Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests->aim->stylegan2_pytorch) (3.6)
  Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests->aim->stylegan2_pytorch) (2.0.7)
  Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests->aim->stylegan2_pytorch) (2024.2.2)
  Requirement already satisfied: mpmath>=0.19 in /usr/local/lib/python3.10/dist-packages (from sympy->torch->stylegan2_pytorch) (1.3.0)
  Requirement already satisfied: pycparser in /usr/local/lib/python3.10/dist-packages (from cffi>=1.12->cryptography>=3.0->aim->stylegan2_pytorch) (2.22)
  Requirement already satisfied: annotated-types>=0.4.0 in /usr/local/lib/python3.10/dist-packages (from pydantic!=1.8,!=1.8.1,!=2.0.0,!=2.0.1,!=2.1.0,<3.0.0,>=1.7.4->fastapi<1,>=0.69.0->aim->stylegan2_pytorch) (0.6.0)
  Requirement already satisfied: pydantic-core==2.16.3 in /usr/local/lib/python3.10/dist-packages (from pydantic!=1.8,!=1.8.1,!=2.0.0,!=2.0.1,!=2.1.0,<3.0.0,>=1.7.4->fastapi<1,>=0.69.0->aim->stylegan2_pytorch) (2.16.3)
  Requirement already satisfied: anyio<5,>=3.4.0 in /usr/local/lib/python3.10/dist-packages (from starlette<0.38.0,>=0.37.2->fastapi<1,>=0.69.0->aim->stylegan2_pytorch) (3.7.1)
  Requirement already satisfied: sniffio>=1.1 in /usr/local/lib/python3.10/dist-packages (from anyio<5,>=3.4.0->starlette<0.38.0,>=0.37.2->fastapi<1,>=0.69.0->aim->stylegan2_pytorch) (1.3.1)
  Requirement already satisfied: exceptiongroup in /usr/local/lib/python3.10/dist-packages (from anyio<5,>=3.4.0->starlette<0.38.0,>=0.37.2->fastapi<1,>=0.69.0->aim->stylegan2_pytorch) (1.2.0)
  `
  
  
  
  However,  as I run these two lines to check whether stylegan2_pytorch package has been successfully installed,
  
  `import site
  print(site.getsitepackages())
  `
  
  it only manage to return only these two packages,
  
  `['/usr/local/lib/python3.10/dist-packages', '/usr/lib/python3/dist-packages', '/usr/lib/python3.10/dist-packages']running import site, print(site.getsitepackages())`
  
  