rewrite

2020-10-11 10:02:33 +08:00
parent 6ea13df465
commit 04c6366c07
24 changed files with 483 additions and 968 deletions
--- a/util/misc.py
+++ b/util/misc.py
@@ -8,7 +8,7 @@ import torch.nn as nn


 def add_spectral_norm(module):
-    if isinstance(module, (nn.Conv2d, nn.ConvTranspose2d)) and not hasattr(module, 'weight_u'):
+    if isinstance(module, (nn.Linear, nn.Conv2d, nn.ConvTranspose2d)) and not hasattr(module, 'weight_u'):
        return nn.utils.spectral_norm(module)
    else:
        return module