Exported loss functions

ppeetteerrs · Mar 22, 2022 · d684ea5 · d684ea5
1 parent b832e12
commit d684ea5
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/stylegan2_torch/__init__.py b/stylegan2_torch/__init__.py
@@ -5,6 +5,7 @@
 from stylegan2_torch.discriminator import Discriminator
 from stylegan2_torch.equalized_lr import Blur, EqualConv2d, EqualLeakyReLU, EqualLinear
 from stylegan2_torch.generator import Generator
+from stylegan2_torch.loss import d_loss, d_reg_loss, g_loss, g_reg_loss
 from stylegan2_torch.utils import Resolution, default_channels
 
 __author__ = "Peter Yuen"
@@ -21,6 +22,10 @@
     "EqualConv2d",
     "EqualLeakyReLU",
     "EqualLinear",
+    "d_loss",
+    "d_reg_loss",
+    "g_loss",
+    "g_reg_loss",
 ]