Merge pull request #25 from MICS-Lab/dev

Dev
MICS-Lab · Nov 17, 2023 · 016f56a · 016f56a
2 parents e936351 + cfc0992
commit 016f56a
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/scyan/model.py b/scyan/model.py
@@ -314,7 +314,7 @@ def on_train_epoch_end(self):
             self.hparams.warm_up is not None
             and self.current_epoch == self.hparams.warm_up[1] - 1
         ):
-            print("Ended warm up epochs")
+            log.info("Ended warm up epochs")
             self.module.prior.prior_std = self.hparams.prior_std
 
     @_requires_fit

diff --git a/scyan/module/distribution.py b/scyan/module/distribution.py
@@ -38,7 +38,8 @@ def prior_std(self):
     @prior_std.setter
     def prior_std(self, std: float) -> None:
         self._prior_std = std
-        self.register_buffer("cov", torch.eye((self.n_markers)) * std**2)
+        cov = torch.eye((self.n_markers)) * std**2
+        self.register_buffer("cov", cov.to(self.device))
         self.normal = distributions.Normal(0, std)
         self.compute_constant_terms()