Migrate all uses of `at::Tensor` to `torch::Tensor`.

[Dan-Flores]

Currently, two files use at::Tensor while the majority of files use torch::Tensor.
For consistency, we should update all uses of at::Tensor to torch::Tensor.