From b6e131e4c1b835dce5b901a1169fa53d2c12b786 Mon Sep 17 00:00:00 2001 From: harisreedhar Date: Tue, 11 Mar 2025 12:59:52 +0530 Subject: [PATCH] changes --- face_swapper/src/training.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/face_swapper/src/training.py b/face_swapper/src/training.py index 7348d8e..11b95a1 100644 --- a/face_swapper/src/training.py +++ b/face_swapper/src/training.py @@ -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()