simplify householder reshape computation + correct test

deel/torchlip/modules/activation.py

@@ -250,9 +250,8 @@ def __init__(self, channels, k_coef_lip: float = 1.0, theta_initializer=None):
                 raise ValueError(f"Unknown initializer {theta_initializer}")
 
     def forward(self, z, axis=1):
-        theta = self.theta.to(z.device).view(1, -1)
-        for _ in range(len(z.shape) - len(theta.shape)):
-            theta = theta.unsqueeze(-1)
+        theta_shape = (1, -1) + (1,) * (len(z.shape) - 2)
+        theta = self.theta.to(z.device).view(theta_shape)
         x, y = z.split(z.shape[axis] // 2, axis)
         selector = (x * torch.sin(0.5 * theta)) - (y * torch.cos(0.5 * theta))
 

tests/test_activations.py

@@ -39,7 +39,7 @@ def check_serialization(layer_type, layer_params):
     m = uft.generate_k_lip_model(layer_type, layer_params, input_shape=(10,), k=1)
     if m is None:
         return
-    optimizer, loss, _ = uft.compile_model(
+    loss, optimizer, _ = uft.compile_model(
         m,
         optimizer=uft.get_instance_framework(uft.SGD, inst_params={"model": m}),
         loss=CategoricalCrossentropy(from_logits=True),