func to apply sn

util/misc.py

@@ -2,6 +2,15 @@ import logging
 from pathlib import Path
 from typing import Optional
 
+import torch.nn as nn
+
+
+def add_spectral_norm(module):
+    if isinstance(module, (nn.Conv2d, nn.ConvTranspose2d)) and not hasattr(module, 'weight_u'):
+        return nn.utils.spectral_norm(module)
+    else:
+        return module
+
 
 def setup_logger(
         name: Optional[str] = None,