pytorch

b1af69f9 - Small refactor for reshape shape inference