Yup, it works pretty well! Here’s the code I use to make it work. Basically, a FireTensor is a TVM compatible object that operates directly on a PyTorch Tensor.
import ctypes
import torch
import tvm

@tvm.register_extension
class FireTensor(object):
    _tvm_tcode = tvm.TypeCode.ARRAY_HANDLE

    def __init__(self, tensor):
        self.handle = torch._C._to_dlpack(tensor)
        self.name = self.get_name()
    def get_name(self):
        ctypes.pythonapi.PyCapsule_GetName.restype = ctypes.c_char_p
        ctypes.pythonapi.PyCapsule_GetName.argtypes = [ctypes.py_object]
        return ctypes.pythonapi.PyCapsule_GetName(self.handle)
    def to_torch(self):
        return torch._C._from_dlpack(self.handle)
    @property
    def _tvm_handle(self):
        ctypes.pythonapi.PyCapsule_GetPointer.restype = ctypes.c_void_p
        ctypes.pythonapi.PyCapsule_GetPointer.argtypes = [ctypes.py_object, ctypes.c_char_p]
        return ctypes.pythonapi.PyCapsule_GetPointer(self.handle, self.name)