Co-authored-by: Brayden Zhong <b8zhong@uwaterloo.ca>
Co-authored-by: Jinyan Chen <jinyanc@nvidia.com>