Update app.py

app.py

@@ -308,11 +308,12 @@ def test(net, testloader):
     accuracy = metric.compute()["accuracy"]
     return loss, accuracy
 
-net = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=2).to(DEVICE)
 
 def main():
     st.write("## Federated Learning with dynamic models and datasets for mobile devices")
     dataset_name = st.selectbox("Dataset", ["imdb", "amazon_polarity", "ag_news"])
+    net = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=2).to(DEVICE)
+
     model_name = st.selectbox("Model", ["facebook/hubert-base-ls960","bert-base-uncased", "distilbert-base-uncased"])
     NUM_CLIENTS = st.slider("Number of Clients", min_value=1, max_value=10, value=2)
     NUM_ROUNDS = st.slider("Number of Rounds", min_value=1, max_value=10, value=3)