Update SettingC-SO2-Spread-B.py

Correct steerer
georg-bn · Jun 30, 2024 · 9c79b60 · 9c79b60
1 parent a51da82
commit 9c79b60
Showing 1 changed file with 20 additions and 14 deletions.
diff --git a/experiments/SettingC-SO2-Spread-B.py b/experiments/SettingC-SO2-Spread-B.py
@@ -76,45 +76,51 @@ def train(detector_weights):
 
     generator = torch.block_diag(
         torch.zeros(
-            [NUM_PROTOTYPES // 2,
-             NUM_PROTOTYPES // 2],
+            [NUM_PROTOTYPES - 12 * (NUM_PROTOTYPES // 14),
+             NUM_PROTOTYPES - 12 * (NUM_PROTOTYPES // 14)],
             device='cuda',
         ),
         *(
             torch.tensor([[0., 1],
                           [-1, 0]],
-                         device='cuda')
-            for _ in range(NUM_PROTOTYPES // 8)
+                         device='cuda',
+                        )
+            for _ in range(NUM_PROTOTYPES // 14)
         ),
         *(
             torch.tensor([[0., 2],
                           [-2, 0]],
-                         device='cuda')
-            for _ in range(NUM_PROTOTYPES // 16)
+                         device='cuda',
+                        )
+            for _ in range(NUM_PROTOTYPES // 14)
         ),
         *(
             torch.tensor([[0., 3],
                           [-3, 0]],
-                         device='cuda')
-            for _ in range(NUM_PROTOTYPES // 32)
+                         device='cuda',
+                        )
+            for _ in range(NUM_PROTOTYPES // 14)
         ),
         *(
             torch.tensor([[0., 4],
                           [-4, 0]],
-                         device='cuda')
-            for _ in range(NUM_PROTOTYPES // 64)
+                         device='cuda',
+                        )
+            for _ in range(NUM_PROTOTYPES // 14)
         ),
         *(
             torch.tensor([[0., 5],
                           [-5, 0]],
-                         device='cuda')
-            for _ in range(NUM_PROTOTYPES // 128)
+                         device='cuda',
+                        )
+            for _ in range(NUM_PROTOTYPES // 14)
         ),
         *(
             torch.tensor([[0., 6],
                           [-6, 0]],
-                         device='cuda')
-            for _ in range(NUM_PROTOTYPES // 128)
+                         device='cuda',
+                        )
+            for _ in range(NUM_PROTOTYPES // 14)
         ),
     )
     steerer = ContinuousSteerer(generator)