Update app.py

app.py

@@ -200,7 +200,7 @@ class MultiMultiSignalingGame:
         return recon_loss + beta * sum(kld_losses), recon_loss, sum(kld_losses)
     
 
-def train_signal_game(NUM_SENDERS, NUM_RECEIVERS, num_rounds):
+def run_signal_game(NUM_SENDERS, NUM_RECEIVERS, num_rounds):
     # para_checker = st.empty()
     # para_checker.text(f"NUM_SENDERS: {NUM_SENDERS}, NUM_RECEIVERS: {NUM_RECEIVERS}, num_rounds: {num_rounds}, EMBEDDING_DIM: {EMBEDDING_DIM}, HIDDEN_DIM: {HIDDEN_DIM}, LATENT_DIM: {LATENT_DIM}, SEQ_LEN: {SEQ_LEN}, TAU: {TAU}, nhead: {NHEAD}, num_layers: {NUM_LAYERS}, BATCH_SIZE: {BATCH_SIZE}")
     senders = [TransformerEncoder().to(device) for _ in range(NUM_SENDERS)]
@@ -246,6 +246,9 @@ def train_signal_game(NUM_SENDERS, NUM_RECEIVERS, num_rounds):
         ax.set_ylabel('Loss')
         ax.legend()
         loss_plot_placeholder.pyplot(fig)
+        # Close the figure to free up memory
+        plt.close(fig)
+        
         progress_bar.progress(round / num_rounds)
         # 刷新显示每次交互的句子
         interaction_str = "\n\n".join([f"Sender {i} -> Receiver {j}\nSend(encode): {input_sentence}\nReceive(decode): {output_sentence}" 
@@ -293,4 +296,4 @@ with advanced_settings:
     BATCH_SIZE = st.slider("BATCH_SIZE", 1, 128, 32)
 
 if st.sidebar.button('Start'):
-    train_signal_game(NUM_SENDERS, NUM_RECEIVERS, num_rounds)
+    run_signal_game(NUM_SENDERS, NUM_RECEIVERS, num_rounds)