set ability to do warmup steps for each unet during training

2025-12-19 09:44:19 +01:00 · 2022-07-05 16:24:16 -07:00
parent 3afdcdfe86
commit ec68243479
4 changed files with 53 additions and 5 deletions
--- a/setup.py
+++ b/setup.py
@@ -37,6 +37,7 @@ setup(
    'packaging',
    'pillow',
    'pydantic',
+    'pytorch-warmup',
    'resize-right>=0.0.2',
    'rotary-embedding-torch',
    'torch>=1.10',