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