update LossRecorder to LossOutput

YosefLab · Jul 11, 2023 · efc5981 · efc5981
1 parent d1934e4
commit efc5981
Showing 1 changed file with 6 additions and 6 deletions.
diff --git a/mrvi/_module.py b/mrvi/_module.py
@@ -2,7 +2,7 @@
 import torch.distributions as db
 import torch.nn as nn
 from scvi import REGISTRY_KEYS
-from scvi.module.base import BaseModuleClass, LossRecorder, auto_move_data
+from scvi.module.base import BaseModuleClass, LossOutput, auto_move_data
 from scvi.nn import one_hot
 from torch.distributions import kl_divergence as kl
 
@@ -206,9 +206,9 @@ def loss(
 
         kl_local = torch.tensor(0.0)
         kl_global = torch.tensor(0.0)
-        return LossRecorder(
-            loss,
-            reconstruction_loss,
-            kl_local,
-            kl_global,
+        return LossOutput(
+            loss = loss,
+            reconstruction_loss = reconstruction_loss,
+            kl_local = kl_local,
+            kl_global = kl_global
         )