diff --git a/codes/models/archs/stylegan/stylegan2.py b/codes/models/archs/stylegan/stylegan2.py
index f29475ce..16735af3 100644
--- a/codes/models/archs/stylegan/stylegan2.py
+++ b/codes/models/archs/stylegan/stylegan2.py
@@ -8,6 +8,7 @@ from random import random
 import torch
 import torch.nn.functional as F
 import models.steps.losses as L
+import numpy as np
 
 from kornia.filters import filter2D
 from linear_attention_transformer import ImageLinearAttention