From 9af049c671c56b17371ab55288dbacf85df58383 Mon Sep 17 00:00:00 2001
From: James Betker <jbetker@gmail.com>
Date: Sat, 14 Nov 2020 22:09:18 -0700
Subject: [PATCH] Import fix for unet

---
 codes/models/archs/stylegan/stylegan2_unet_disc.py | 3 +++
 1 file changed, 3 insertions(+)

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)