Minor bugfix for c2st

[psteinb]

Just by chance: how backwards compatible is this?

[michaeldeistler]

Both axis and dim work for torch.mean. But dim is the "usual" way of doing it, see here