From f3e1c3feaa7a95886518da6b5b2844212e7c9d72 Mon Sep 17 00:00:00 2001 From: henryruhs Date: Wed, 15 Jan 2025 09:26:17 +0100 Subject: [PATCH] Improve naming for arcface converter trainer --- arcface_converter/src/training.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/arcface_converter/src/training.py b/arcface_converter/src/training.py index 7cde3c5..880a717 100644 --- a/arcface_converter/src/training.py +++ b/arcface_converter/src/training.py @@ -30,16 +30,16 @@ class ArcFaceConverterTrainer(pytorch_lightning.LightningModule): return self.model(source_embedding) def training_step(self, batch : Batch, batch_index : int) -> Tensor: - source, target = batch - output = self(source) - loss = self.loss_fn(output, target) + source_embedding, target_embedding = batch + output_embedding = self(source_embedding) + loss = self.loss_fn(output_embedding, target_embedding) self.log('train_loss', loss, prog_bar = True, logger = True) return loss def validation_step(self, batch : Batch, batch_index : int) -> Tensor: - source, target = batch - output = self(source) - loss = self.loss_fn(output, target) + source_embedding, target_embedding = batch + output_embedding = self(source_embedding) + loss = self.loss_fn(output_embedding, target_embedding) self.log('val_loss', loss, prog_bar = True, logger = True) return loss