diff options
| -rw-r--r-- | text_recognizer/models/vq_transformer.py | 6 |
1 files changed, 2 insertions, 4 deletions
diff --git a/text_recognizer/models/vq_transformer.py b/text_recognizer/models/vq_transformer.py index c07efa4..8ec28fd 100644 --- a/text_recognizer/models/vq_transformer.py +++ b/text_recognizer/models/vq_transformer.py @@ -12,8 +12,6 @@ from text_recognizer.models.transformer import TransformerLitModel class VqTransformerLitModel(TransformerLitModel): """A PyTorch Lightning model for transformer networks.""" - alpha: float = attr.ib(default=1.0) - def forward(self, data: Tensor) -> Tensor: """Forward pass with the transformer network.""" return self.predict(data) @@ -22,7 +20,7 @@ class VqTransformerLitModel(TransformerLitModel): """Training step.""" data, targets = batch logits, commitment_loss = self.network(data, targets[:, :-1]) - loss = self.loss_fn(logits, targets[:, 1:]) + self.alpha * commitment_loss + loss = self.loss_fn(logits, targets[:, 1:]) + commitment_loss self.log("train/loss", loss) self.log("train/commitment_loss", commitment_loss) return loss @@ -31,7 +29,7 @@ class VqTransformerLitModel(TransformerLitModel): """Validation step.""" data, targets = batch logits, commitment_loss = self.network(data, targets[:, :-1]) - loss = self.loss_fn(logits, targets[:, 1:]) + self.alpha * commitment_loss + loss = self.loss_fn(logits, targets[:, 1:]) + commitment_loss self.log("val/loss", loss, prog_bar=True) self.log("val/commitment_loss", commitment_loss) |