Merge pull request NVlabs#3 from pbizimis/main

Add weight layer to cff
snakch · Mar 24, 2021 · 182a62c · 182a62c
2 parents 6e64aa0 + d24f806
commit 182a62c
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/closed_form_factorization.py b/closed_form_factorization.py
@@ -21,7 +21,7 @@
     modulate = {
         k[0]: k[1]
         for k in G.named_parameters()
-        if "affine" in k[0] and "torgb" not in k[0] and "weight" in k[0]
+        if "affine" in k[0] and "torgb" not in k[0] and "weight" in k[0] or ("torgb" in k[0] and "b4" in k[0] and "weight" in k[0] and "affine" in k[0])
     }
 
     weight_mat = []