Add ShapeInference for AtomicIter Op (#20021)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/20021
Add shape inference for AtomicIter operator. The operator takes two blobs iteration and iter_mutex as input and outputs iteration, which should have the same type and shape as the input.
Reviewed By: un-disclosed
Differential Revision: D15111643
fbshipit-source-id: 0d06413305cc4c6257c0cfabf62fb874970803bc
