First thanks for your work! when I run your example script "The setup of the pFedSD for simple cnn on cifar10 with dirichlet distribution", and a bug accurs as the image below shows. i search on the stackoverflow and someone said decreasing the batch size may works. However when I change the batch size from 64 to 16, same bug accurs.
The code is running on 2 NVIDIA 3090 GPUs.

First thanks for your work! when I run your example script "The setup of the pFedSD for simple cnn on cifar10 with dirichlet distribution", and a bug accurs as the image below shows. i search on the stackoverflow and someone said decreasing the batch size may works. However when I change the batch size from 64 to 16, same bug accurs.

The code is running on 2 NVIDIA 3090 GPUs.