Update app.py

app.py

@@ -225,7 +225,11 @@ def beam_generate_v2(model,src_tensor, beam=5, max_len=50, alpha=0.7):
 
                 # decoder step: input is last token
                 dec_in = torch.tensor([[seq[-1]]], device=DEVICE)
-                out, new_h = model.decoder(dec_in, hid, enc_out)
+                # Call decoder with correct arguments depending on model type
+                if isinstance(model.decoder, Decoder_with_attn):
+                    out, new_h = model.decoder(dec_in, hid, enc_out)
+                else:
+                    out, new_h = model.decoder(dec_in, hid)
                 probs = F.log_softmax(out, dim=-1).squeeze(0)  # [vocab]
 
                 # penalty for repetition