pytorch - 💡(How to fix) Fix future annotations break torch.library.triton_op

repro:

from __future__ import annotations

from torch.library import triton_op


def register():
    import torch

    @triton_op("repro::bad", mutates_args=())
    def bad(x: torch.Tensor) -> torch.Tensor:
        return x

register()

ValueError: infer_schema(func): Unsupported type annotation torch.Tensor. It is not a type. Got func with signature (x: 'torch.Tensor') -> 'torch.Tensor')

Also, this error is difficult to diagnose, because there is a try-finally that just swallows the following: "NameError: name 'torch' is not defined". Would love it for this one to be fixed too.

meta-internal context

cc @lolpack @maggiemoss @ndmitchell @kinto0 @samwgoldman @anjali411 @chauhang @penguinwu @bdhirsh @bobrenjc93 @aorenste