fix typo

2024-05-24 20:04:14 -04:00 · 2024-05-24 20:04:14 -04:00 · 5f6b66562f
parent 547d3c3ea1
commit 5f6b66562f
1 changed files with 1 additions and 1 deletions
--- a/lerobot/common/policies/vqbet/modeling_vqbet.py
+++ b/lerobot/common/policies/vqbet/modeling_vqbet.py
@ -801,7 +801,7 @@ class VqVae(nn.Module):
        dec_out = self.decoder(state_vq)
        encoder_loss = (state - dec_out).abs().mean()

-        rep_loss = encoder_loss * vq_loss_state * 5
+        rep_loss = encoder_loss + vq_loss_state * 5

        metric = (
            encoder_loss.clone().detach(),