Update app.py

app.py

@@ -234,7 +234,7 @@ import streamlit
 
 # If you're curious of all the loggers
 
-DEVICE = torch.device("cpu")
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 fl.common.logger.configure(identifier="myFlowerExperiment", filename="./log.txt")
 
 def load_data(dataset_name, train_size=20, test_size=20, num_clients=2):
@@ -510,8 +510,8 @@ def main():
                 num_clients=NUM_CLIENTS,
                 config=fl.server.ServerConfig(num_rounds=1),
                 strategy=strategy,
-                client_resources={"num_cpus": 1, "num_gpus": 0},
-                ray_init_args={"log_to_driver": True, "num_cpus": 1, "num_gpus": 0}
+                client_resources={"num_cpus": 1, "num_gpus": (1 if torch.cuda.is_available() else 0)},
+                ray_init_args={"log_to_driver": True, "num_cpus": 1, "num_gpus": (1 if torch.cuda.is_available() else 0)}
             )
 
             for i, client in enumerate(clients):