How to Implement Shuffle BN for MoCo in Tensorflow 2

Hi, hope I can get some help here. I want to implement unsupervised contrastive learning model MoCo in TF2, but I have no idea how to implement the essential trick mentioned in the paper - Shuffling BN. I think I understand what shuffling BN does, but I don’t know any APIs to fetch different data slices from each GPU, shuffle them, and send new slices back to each GPU.