CPU usage reaches 100%, while the GPU (NVIDIA GeForce RTX 3060) only uses 1.2GB of memory.

[YANGTUOMAO]

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!