All_gather get stuck when using attention_mask having 0

distributed
2

I gather all_predictions and all_labels using all_gather_object() instead of gather batch_label and batch_pred. And the stuck problem is solved. However, I still cannot find the cause of getting stuck