pytorch

3b188c58 - Don't use subclass when tracing and call wait_tensor immediately. (#98001)

Don't use subclass when tracing and call wait_tensor immediately. (#98001) This change expects that proper scheduling of the wait_tensor call will happen over the traced graph. Pull Request resolved: https://github.com/pytorch/pytorch/pull/98001 Approved by: https://github.com/wconstab, https://github.com/wanchaol