Hi,
The faulty Tensor is of size (128,), so I don’t think the issue is with the output but more with the weight or bias of the batchnorm no?
Do you use an optimizer or update these weights/bias between the forward and backward?