Format code with black

jameschapman19 · Sep 19, 2023 · 50f693d · 50f693d
1 parent 5d482ef
commit 50f693d
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 7 deletions.
diff --git a/cca_zoo/deep/callbacks.py b/cca_zoo/deep/callbacks.py
@@ -23,22 +23,42 @@ def on_train_epoch_end(self, trainer: Trainer, pl_module: LightningModule) -> No
             train_corr,
         )
 
+
 class MinibatchTrainCorrelationCallback(Callback):
-    mcca=MCCA()
-    def on_train_batch_end(self, trainer: Trainer, pl_module: LightningModule, outputs, batch, batch_idx, dataloader_idx):
+    mcca = MCCA()
+
+    def on_train_batch_end(
+        self,
+        trainer: Trainer,
+        pl_module: LightningModule,
+        outputs,
+        batch,
+        batch_idx,
+        dataloader_idx,
+    ):
         with torch.no_grad():
             train_corr = self.mcca.loss(pl_module(batch["views"])).sum()
             pl_module.log(
                 "train/corr",
                 train_corr,
             )
 
+
 class MinibatchValidationCorrelationCallback(Callback):
-    mcca=MCCA()
-    def on_validation_batch_end(self, trainer: Trainer, pl_module: LightningModule, outputs, batch, batch_idx, dataloader_idx):
+    mcca = MCCA()
+
+    def on_validation_batch_end(
+        self,
+        trainer: Trainer,
+        pl_module: LightningModule,
+        outputs,
+        batch,
+        batch_idx,
+        dataloader_idx,
+    ):
         with torch.no_grad():
             val_corr = self.mcca.loss(pl_module(batch["views"])).sum()
             pl_module.log(
                 "val/corr",
                 val_corr,
-            )
+            )
diff --git a/cca_zoo/deep/objectives.py b/cca_zoo/deep/objectives.py
@@ -140,8 +140,8 @@ def correlation(self, views):
         views = _demean(views)
 
         SigmaHat12 = torch.cov(torch.hstack((views[0], views[1])).T)[
-                     : latent_dims, latent_dims :
-                     ]
+            :latent_dims, latent_dims:
+        ]
         SigmaHat11 = torch.cov(views[0].T) + self.r * torch.eye(
             o1, device=views[0].device
         )