Update app.py

app.py

@@ -234,8 +234,12 @@ def train_signal_game(NUM_SENDERS, NUM_RECEIVERS, num_rounds):
     loss_plot_placeholder = st.empty()  # 创建一个空位占位符来显示损失图
     interactions_placeholder = st.empty()  # 创建一个空位占位符来显示交互
 
-    for round in range(num_rounds):
-        states = [torch.randint(VOCAB_SIZE, (BATCH_SIZE, 16)).to(device) for _ in range(NUM_SENDERS)]
+    for round, batch in enumerate(train_dataloader):
+        if round >= num_rounds:
+            break
+        states = [batch.to(device) for _ in range(NUM_SENDERS)]
+    # for round in range(num_rounds):
+        # states = [torch.randint(VOCAB_SIZE, (BATCH_SIZE, 16)).to(device) for _ in range(NUM_SENDERS)]
         loss, recon_loss, kld_loss, interactions = game.play_round(states)
         losses.append(loss)
         recon_losses.append(recon_loss)