[TensorExpr] Fix the shape info check in the TE fuser pass. (#37882)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/37882
Previously we checked if a node's inputs and outputs have shape
info only when we tried to merge this node into an existing fusion
group, but we didn't check it for the first node in the group. This PR
fixes that. It was causing a failure on test_milstm_cuda, which is now
fixed.
Test Plan: Imported from OSS
Reviewed By: Krovatkin
Differential Revision: D21412756
Pulled By: ZolotukhinM
fbshipit-source-id: 3ca30637ab8fe68443adb5fc03f1b8a11085a6a8
Author
Mikhail Zolotukhin
