

Issue Description
I have noticed that when running the program, the CPU usage reaches 100%, while the GPU (NVIDIA GeForce RTX 3060) only uses 1.2GB of memory. Despite specifying the device as cuda:0 in the command, it seems that the program is still utilizing the CPU for computation.
Environment Information
Operating System: Windows 11
Python Version: 3.x (please replace with the actual version)
Running Environment: Using a virtual environment (e.g., Anaconda)
Solutions Attempted
Ensured that the device is correctly specified as cuda:0.
Verified that CUDA and PyTorch are properly installed.
I am a beginner in federated learning and would appreciate any further suggestions or possible solutions. Thank you!