Merge remote-tracking branch 'origin/main'

cca_zoo/linear/_gradient/_base.py

@@ -93,7 +93,12 @@ def _fit(self, views: Iterable[np.ndarray]):
         )
         if self.batch_size is None:
             # if the batch size is None, put views on the device
-            self.batch={"views":[view.to(trainer._accelerator_connector._accelerator_flag) for view in train_dataset.views]}
+            self.batch = {
+                "views": [
+                    view.to(trainer._accelerator_connector._accelerator_flag)
+                    for view in train_dataset.views
+                ]
+            }
         trainer.fit(self, train_dataloader, val_dataloader)
         # return the weights from the module. They will need to be changed from torch tensors to numpy arrays
         weights = [weight.detach().cpu().numpy() for weight in self.torch_weights]

cca_zoo/linear/_gradient/_ey.py

@@ -38,7 +38,9 @@ def _compute_loss(self, batch) -> dict:
         else:
             if len(self.batch_queue) < 5:
                 self.batch_queue.append(batch)
-                return {"loss": torch.tensor(0, requires_grad=True, dtype=torch.float32)}
+                return {
+                    "loss": torch.tensor(0, requires_grad=True, dtype=torch.float32)
+                }
             else:
                 random_batch = self._get_random_batch()
                 loss = self.loss(batch["views"], random_batch["views"])

cca_zoo/linear/_gradient/_stochasticpls.py

@@ -10,7 +10,7 @@ def _more_tags(self):
 
     def training_step(self, batch, batch_idx):
         if batch is None:
-            batch=dict(("views",self.data))
+            batch = dict(("views", self.data))
         for weight in self.torch_weights:
             weight.data = self._orth(weight)
         scores = self(batch["views"])