Generalize some TensorIterator consumers to take TensorIteratorBase (#48727)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/48727
Signed-off-by: Edward Z. Yang <ezyang@fb.com>
Test Plan: Imported from OSS
Reviewed By: bhosmer
Differential Revision: D25278033
Pulled By: ezyang
fbshipit-source-id: 77f125ddb8446edf467a22130227d90583884bca
