wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-01-11 06:40:48 +08:00
Code Issues Actions 12 Packages Projects Releases Wiki Activity

Fix no-member errors in comfy.ldm.modules.ema

Browse Source
This commit is contained in:
Max Tretikov 2024-06-14 09:01:38 -06:00
parent f0812a88fe
commit 05f4c2a5bc
1 changed files with 1 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
comfy/ldm/modules/ema.py
View File

@ -23,14 +23,13 @@ class LitEma(nn.Module):
self.collected_params = [] self.collected_params = []
def reset_num_updates(self): def reset_num_updates(self):
del self.num_updates
self.register_buffer('num_updates', torch.tensor(0, dtype=torch.int)) self.register_buffer('num_updates', torch.tensor(0, dtype=torch.int))
def forward(self, model): def forward(self, model):
decay = self.decay decay = self.decay
if self.num_updates >= 0: if self.num_updates >= 0:
self.num_updates += 1 self.register_buffer('num_updates', torch.tensor(1 + self.num_updates, dtype=torch.int))
decay = min(self.decay, (1 + self.num_updates) / (10 + self.num_updates)) decay = min(self.decay, (1 + self.num_updates) / (10 + self.num_updates))
one_minus_decay = 1.0 - decay one_minus_decay = 1.0 - decay