Looking at your code some more, I wonder whether it might be more reasonable to take the mean over the dimension(s >=) 1 and then do a scatter_add_ and divide for handling the 0 dimension. I don’t want to keep you from doing your own kernel, but I would guess that it’s hard to win much over that.
Accessors only work when you know the dimension (i.e. how many “axes”) of the tensor at compile time. You very likely want at least the 1.0 preview, and you need the updates from the batch norm PR for good cuda tensor accessors (e.g. you want 32bit indexing if you can use it etc.).
Best regards
Thomas