Repository: pytorch
Issue: PyTorch crashes without an error message, when running this code snippet with torch.tensor subclassing & forward hooks (Not sure what the exact cause is, but the code snippet reliably causes it)
Labels: triaged,Stale,tensor subclass

### 🐛 Describe the bug

While working a project for PyTorch's [Captum](https://github.com/pytorch/captum) library, I came across a bug that I've been struggling to narrow down the cause of. I've done my best to simplify what is happening in the Captum code, and the snippet of code below should reliably reproduce the crash, though I apologies for not being able to narrow it down to smaller snippet of code.

The example code uses torch.tensor subclassing and forward hooks, and they appear to be important for causing the crash.

I have no idea if there should be an error message when running the code, or if there should be no issue at all.

```
import torch
from torchvision import models
model = models.resnet18()

from typing import Type, Callable, List, Tuple, Union
import numpy as np
from types import MethodType


class TestTensor(torch.Tensor):
    @staticmethod
    def __new__(
        cls: Type["TestTensor"],
        x: Union[List, np.ndarray, torch.Tensor] = [],
        *args,
        **kwargs,
    ) -> torch.Tensor:
        if isinstance(x, torch.Tensor) and x.is_cuda:
            x.show = MethodType(cls.show, x)
            x.export = MethodType(cls.export, x)
            return x
        else:
            return super().__new__(cls, x, *args, **kwargs)

    @classmethod
    def __torch_function__(
        cls: Type["TestTensor"],
        func: Callable,
        types: List[Type[torch.Tensor]],
        args: Tuple = (),
        kwargs: dict = None,
    ) -> torch.Tensor:
        if kwargs is None:
            kwargs = {}
        return super().__torch_function__(func, types, args, kwargs)


class TestTensor2(torch.nn.Module):

    def __init__(self):
        super().__init__()
        self.test_tensor = torch.randn(3,3,224,224).clamp(0,1)

    def forward(self):
        x = self.test_tensor
        return TestTensor(x)


def test_hook(target):

    def forward_hook(self, input, output) -> None:
         pass

    test_hooks = target.register_forward_hook(forward_hook)
    test_hooks.remove()
    return image().detach(), torch.randn(5)


class CaptumModuleOutputsHook:
    def __init__(self, target_modules) -> None:
        self.outputs = dict.fromkeys(target_modules, None)
        self.hooks = [
            module.register_forward_hook(self._forward_hook())
            for module in target_modules
        ]

    def _forward_hook(self) -> Callable:
        def forward_hook(
            module: torch.nn.Module, input: Tuple[torch.Tensor], output: torch.Tensor
        ) -> None:
            assert module in self.outputs.keys()
            self.outputs[module] = output

        return forward_hook

    def consume_outputs(self):
        outputs = self.outputs
        self.outputs = dict.fromkeys(self.outputs.keys(), None)
        return outputs

    def remove_hooks(self) -> None:
        for hook in self.hooks:
            hook.remove()


def collect_activations(model, target, input_tensor):
    layers = CaptumModuleOutputsHook(target)
    try:
        model(input_tensor)
        activations_dict = layers.consume_outputs()
    finally:
        layers.remove_hooks()
    return activations_dict[target[0]]


def trigger_crash(
    model,
    image,
    target,
):
    attempts, attempt_losses = [], []

    # Removing this loop somehow prevents the crash from happening
    for a in range(1):
        imgs, losses = test_hook(target)
        attempts.append(imgs.detach()); attempt_losses.append(losses)
    final_image, final_losses = torch.cat(attempts, 0), torch.stack(attempt_losses)

    activ = collect_activations(model, [target], final_image) # Crash happens on this line

    # Commenting out these lines of code somehow prevents the crash from happening
    comparison_losses = torch.stack([activ.mean()]*3)
    sorted_idx = torch.sort(comparison_losses)[1]
    best_image = final_image[sorted_idx[0:3]]
    best_losses = final_losses[sorted_idx[0:3]]
    return best_image, best_losses


image = TestTensor2()
trigger_crash(model, image, model.layer1)
```

### Versions

```
PyTorch version: 1.10.0+cu111
Is debug build: False
CUDA used to build PyTorch: 11.1
ROCM used to build PyTorch: N/A

OS: Ubuntu 18.04.5 LTS (x86_64)
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
Clang version: 6.0.0-1ubuntu2 (tags/RELEASE_600/final)
CMake version: version 3.12.0
Libc version: glibc-2.26

Python version: 3.7.12 (default, Sep 10 2021, 00:21:48)  [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-5.4.144+-x86_64-with-Ubuntu-18.04-bionic
Is CUDA available: False
CUDA runtime version: 11.1.105
GPU models and configuration: Could not collect
Nvidia driver version: Could not collect
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.7.6.5
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.0.5
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.19.5
[pip3] torch==1.10.0+cu111
[pip3] torchaudio==0.10.0+cu111
[pip3] torchsummary==1.5.1
[pip3] torchtext==0.11.0
[pip3] torchvision==0.11.1+cu111
[conda] Could not collect
```

Basically a standard Google Colab instance, but the issue should also occur on Windows 10 as well.