DEF for fake vae and adjust scaling

2025-09-12 18:09:08 -06:00
parent b95c17dc17
commit 3666b112a8
2 changed files with 5 additions and 4 deletions
--- a/toolkit/models/FakeVAE.py
+++ b/toolkit/models/FakeVAE.py
@@ -25,11 +25,12 @@ class Config:


 class FakeVAE(nn.Module):
-    def __init__(self):
+    def __init__(self, scaling_factor=1.0):
        super().__init__()
        self._dtype = torch.float32
        self._device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.config = Config()
+        self.config.scaling_factor = scaling_factor

    @property
    def dtype(self):