CalvinBackup/facefusion-labs
1
0
Fork 0
mirror of https://github.com/facefusion/facefusion-labs.git synced 2026-06-25 07:59:55 +02:00
Code Issues Packages Projects Releases Wiki Activity

changes

Browse Source
This commit is contained in:
harisreedhar
2025-03-11 12:59:52 +05:30
committed by henryruhs
parent d809f66216
commit b6e131e4c1
1 changed files with 4 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
face_swapper/src/training.py
+4 -4
View File
@@ -96,7 +96,7 @@ class FaceSwapperTrainer(LightningModule):
def training_step(self, batch : Batch, batch_index : int) -> Tensor:
source_tensor, target_tensor = batch
is_accumulation_step = (batch_index + 1) % self.config_accumulate_size == 0
do_update = (batch_index + 1) % self.config_accumulate_size == 0
generator_optimizer, discriminator_optimizer, masker_optimizer = self.optimizers() #type:ignore[attr-defined]
source_embedding = calc_embedding(self.embedder, source_tensor, (0, 0, 0, 0))
target_attributes = self.generator.get_attributes(target_tensor)
@@ -115,7 +115,7 @@ class FaceSwapperTrainer(LightningModule):
self.manual_backward(generator_loss)
if is_accumulation_step:
if do_update:
generator_optimizer.step()
generator_optimizer.zero_grad()
@@ -127,7 +127,7 @@ class FaceSwapperTrainer(LightningModule):
self.manual_backward(mask_loss)
if is_accumulation_step:
if do_update:
masker_optimizer.step()
masker_optimizer.zero_grad()
@@ -140,7 +140,7 @@ class FaceSwapperTrainer(LightningModule):
self.manual_backward(discriminator_loss)
if is_accumulation_step:
if do_update:
discriminator_optimizer.step()
discriminator_optimizer.zero_grad()