**Is your feature request related to a problem? Please describe.** check compatibility with the latest pytorch