Add a 'dim' argument to nuclear norm (#21022)
Summary:
Addresses #18275.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/21022
Differential Revision: D15743515
Pulled By: ezyang
fbshipit-source-id: e4aaea0bd7f863a2abad45c4322d6a9fb02a88e3
