diff --git a/codes/models/archs/stylegan/stylegan2_unet_disc.py b/codes/models/archs/stylegan/stylegan2_unet_disc.py
index fd51e1ea..634095b6 100644
--- a/codes/models/archs/stylegan/stylegan2_unet_disc.py
+++ b/codes/models/archs/stylegan/stylegan2_unet_disc.py
@@ -4,6 +4,9 @@ from math import log2
 import torch
 import torch.nn as nn
 
+from models.archs.stylegan.stylegan2 import attn_and_ff
+
+
 def leaky_relu(p=0.2):
     return nn.LeakyReLU(p)