[pt] Replace size(dim) with sizes()[dim] (#49255)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/49255
- Replace `size(dim)` with `sizes()[dim]` because `sizes()` does not go through the dispatcher and is marginally better.
- Remove unnecessary `size(dim)` and `sizes()` calls by saving the return value of `sizes()` to a temporary var.
Reviewed By: radkris-git
Differential Revision: D25488129
fbshipit-source-id: 4039e0609df20d5888666a71ad93b15e9a2182c5
