Reductions are sensitive to overflow if you are using FP16.
You should perform all reductions in FP32 just to make sure to get a valid result.
Just call .float() on your tensor before passing it to torch.mean().
Autograd will rewind this operation in the backward pass so that your model will still be in half precision.