Update app.py

app.py

@@ -234,12 +234,8 @@ def train_signal_game(NUM_SENDERS, NUM_RECEIVERS, num_rounds):
     loss_plot_placeholder = st.empty()  # 创建一个空位占位符来显示损失图
     interactions_placeholder = st.empty()  # 创建一个空位占位符来显示交互
 
-    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)]
+    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)