Browse Source

allow lists or tuples for weights as well

Jens Keim 3 years ago
parent
commit
a39ff65a6a
1 changed files with 4 additions and 1 deletions
  1. 4 1
      privacy_engine_xl.py

+ 4 - 1
privacy_engine_xl.py

@@ -70,7 +70,10 @@ def apply_noise(weights, batch_size, max_norm, noise_multiplier, noise_type, dev
         The method of loss reduction.
         currently supported: mean
     """
-    for p in weights.values():
+    if isinstance(weights, dict):
+        weights = weights.values()
+
+    for p in weights:
         noise = generate_noise(max_norm, p, noise_multiplier, noise_type, device)
         if loss_reduction == "mean":
             noise /= batch_size