Fix returning instance of custom class from method
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/32312
Test Plan: Imported from OSS
Differential Revision: D19433511
Pulled By: jamesr66a
fbshipit-source-id: f048d5f60eaba992ee42fea2d318a59b3a156578
