import mxnet as mx
from mxnet import gluon, autograd
from mxnet.gluon import nn
from mxnet.gluon.nn import Block
# import SyncBN here
from syncbn import BatchNorm, ModelDataParallel

# create your own Block
class Net(Block):
    def __init__(self):
        super(Net, self).__init__()
        self.conv = nn.Conv2D(in_channels=3, channels=10,
                              kernel_size=3, padding=1)
        self.bn = BatchNorm(in_channels=10)
        self.relu = nn.Activation('relu')

    def forward(self, x):
        x = self.conv(x)
        x = self.bn(x)
        x = self.relu(x)
        return x

# set the contexts (suppose using 4 GPUs)
nGPUs = 4
ctx_list = [mx.gpu(i) for i in range(nGPUs)]
# get the model
model = Net()
model.initialize()
model = ModelDataParallel(model, ctx_list)
# load the data
data = mx.random.uniform(-1,1,(8, 3, 24, 24))
x = gluon.utils.split_and_load(data, ctx_list=ctx_list)
with autograd.record():
    y = model(x)