fix logits dropout bug

rwth-i6 · Oct 15, 2024 · 6a147b7 · 6a147b7
1 parent dcf0381
commit 6a147b7
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/i6_models/decoder/attention.py b/i6_models/decoder/attention.py
@@ -201,8 +201,8 @@ def forward(
         readout_in = readout_in.view(readout_in.size(0), readout_in.size(1), -1, 2)  # [B,N,D/2,2]
         readout, _ = torch.max(readout_in, dim=-1)  # [B,N,D/2]
 
-        output = self.output(readout)
-        decoder_logits = self.output_dropout(output)
+        readout_drop = self.output_dropout(readout)
+        decoder_logits = self.output(readout_drop)
 
         state = lstm_state, att_context, accum_att_weights