Thanks for your response. Do you have any similiar thing by using which I can get the input and output for batchnorm layers?