Solved by ProGamerGov in post #2 
              
              
                After some testing, I think this is how you properly type hint __torch_function__: 
from typing import Any, Callable, Dict, Optional, Tuple

class LoggingTensor(torch.Tensor):

    @classmethod
    def __torch_function__(cls, func: Callable, types: Tuple, args: Tuple = (), kwargs: Optional[Dict] = N…