added beta_schedule argument

2025-12-20 02:04:19 +01:00 · 2022-04-17 15:19:33 +02:00
parent 42d6e47387
commit 51361c2d15
1 changed files with 4 additions and 3 deletions
--- a/dalle2_pytorch/dalle2_pytorch.py
+++ b/dalle2_pytorch/dalle2_pytorch.py
@@ -971,7 +971,8 @@ class Decoder(nn.Module):
        clip,
        timesteps=1000,
        cond_drop_prob=0.2,
-        loss_type = 'l1'
+        loss_type="l1",
+        beta_schedule="cosine",
    ):
        super().__init__()
        assert isinstance(clip, CLIP)