Repository: pytorch
Issue: Bug when using `nn.Linear`
Labels: needs reproduction,module: cuda,triaged,module: linear algebra

### 🐛 Describe the bug

I have the strangest bug that I will try my best to explain. Note that the error only occurs when I use the GPU, and when using CPU works as expected with no errors. Note that the error that kills the script does not come from this line, but it is the cause of it -- the errors flags when I try to take a softmax of these nan values and passing into a Categorical distribution.

When using `torch.nn.Linear()` in my network, the first time I use it I get a tensor with nan's. I looked into this thoroughly -- the input has no nan values and the weights of the linear layer also have no nan values (in fact, they have not been updated at all at this point).


The strangest part is that if I go into debug mode in PyCharm and run this line in the terminal and then run it a second time, the first gives me the tensor of nan's and the second will give me the expected output. Assume my input is a tensor `x1`, then what I do is run the line `self.input_transform(x1)` twice (where `self.input_transform = nn.Linear(x, y)`) with the first output being a tensor of nan's and the second being as expected; note that we checked `x1` to make sure nothing was changing in place between the first and second call of `self.input_transform`. 

I have worked with a colleague to try everything to catch this bug: this included doing the matrix multiplication manually with `torch.mm`, changing the input to be `torch.ones_like`, and similar with the data for the weights and bias of the linear layer. No combination of these worked, but if I include `_ = self.input_transform = nn.Linear(x1)` before `x1 = self.input_transform(x1)` then everything flows smoothly. I'll attach the code below for the full network; note that it uses some layers from Pytorch Geometric but this is not causing the issue AFAIK as the linear layer is the first line in the network. 

The error is the first line in the `get_node_embeddings` method. As I said, we have checked to ensure there are no nan values in the input `x1` and have even tried with `torch.one_like` as forced input in the network. 

```
class GMN(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super(GMN, self).__init__()

        self.node_input_transform = nn.Linear(input_dim, hidden_dim)
        self.edge_input_transform = nn.Linear(1, hidden_dim)
        self.mlp_gate = nn.Linear(hidden_dim, hidden_dim)
        self.mlp = nn.Linear(hidden_dim, hidden_dim)
        self.mp1 = GMLayer(hidden_dim)
        self.mp2 = GMLayer(hidden_dim)
        self.mp3 = GMLayer(hidden_dim)
        self.mp4 = GMLayer(hidden_dim)
        self.mp5 = GMLayer(hidden_dim)

        # pooling layers
        self.pooling = nn.Linear(hidden_dim, hidden_dim)
        self.output_layer = nn.Linear(hidden_dim, output_dim)

        # softmax layer
        self.softmax = torch.nn.Softmax(dim=1)

    def forward(self, graph_1, graph_2, node_embds=False):
        x1, edge_index1, edge_attr1, batch1 = graph_1.x, graph_1.edge_index, graph_1.edge_attr, graph_1.batch  # this essentially gets the relevant information from PyG Data objects
        x2, edge_index2, edge_attr2, batch2 = graph_2.x, graph_2.edge_index, graph_2.edge_attr, graph_2.batch

        x1, x2 = self.get_node_embeddings(x1, edge_index1, edge_attr1, x2, edge_index2, edge_attr2)
        if node_embds:
            return x1

        g1_embedding, g2_embedding = self.get_graph_embedding(x1, batch1), self.get_graph_embedding(x2, batch2)

        return torch.sum((g1_embedding - g2_embedding) ** 2, dim=1).view(g1_embedding.size(0), 1), g1_embedding, g2_embedding

    def get_node_embeddings(self, x1, edge_index1, edge_attr1, x2, edge_index2, edge_attr2):
        x1 = self.node_input_transform(x1)  # this line here causes the error
        x2 = self.node_input_transform(x2)
        edge_attr1 = self.edge_input_transform(edge_attr1)
        edge_attr2 = self.edge_input_transform(edge_attr2)

        mv1 = self.get_matching_vector(x1, x2)
        mv2 = self.get_matching_vector(x2, x1)
        x1 = self.mp1.forward(edge_index1, x1, mv1, edge_attr1)
        x2 = self.mp1.forward(edge_index2, x2, mv2, edge_attr2)

        mv1 = self.get_matching_vector(x1, x2)
        mv2 = self.get_matching_vector(x2, x1)
        x1 = self.mp2.forward(edge_index1, x1, mv1, edge_attr1)
        x2 = self.mp2.forward(edge_index2, x2, mv2, edge_attr2)

        mv1 = self.get_matching_vector(x1, x2)
        mv2 = self.get_matching_vector(x2, x1)
        x1 = self.mp3.forward(edge_index1, x1, mv1, edge_attr1)
        x2 = self.mp3.forward(edge_index2, x2, mv2, edge_attr2)

        mv1 = self.get_matching_vector(x1, x2)
        mv2 = self.get_matching_vector(x2, x1)
        x1 = self.mp4.forward(edge_index1, x1, mv1, edge_attr1)
        x2 = self.mp4.forward(edge_index2, x2, mv2, edge_attr2)

        mv1 = self.get_matching_vector(x1, x2)
        mv2 = self.get_matching_vector(x2, x1)
        x1 = self.mp5.forward(edge_index1, x1, mv1, edge_attr1)
        x2 = self.mp5.forward(edge_index2, x2, mv2, edge_attr2)

        return x1, x2
```

### Versions

```
Collecting environment information...
PyTorch version: 1.10.2
Is debug build: False
CUDA used to build PyTorch: 10.2
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.10.2
Libc version: glibc-2.27

Python version: 3.9.7 (default, Sep 16 2021, 13:09:58)  [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-4.15.0-169-generic-x86_64-with-glibc2.27
Is CUDA available: True
CUDA runtime version: 9.1.85
GPU models and configuration: GPU 0: NVIDIA TITAN V
Nvidia driver version: 470.103.01
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] numpy==1.21.2
[pip3] torch==1.10.2
[pip3] torch-cluster==1.5.9
[pip3] torch-geometric==2.0.4
[pip3] torch-scatter==2.0.9
[pip3] torch-sparse==0.6.12
[pip3] torch-spline-conv==1.2.1
[pip3] torchaudio==0.10.2
[pip3] torchvision==0.11.3
[conda] blas                      1.0                         mkl  
[conda] cudatoolkit               10.2.89              hfd86e86_1  
[conda] ffmpeg                    4.3                  hf484d3e_0    pytorch
[conda] mkl                       2021.4.0           h06a4308_640  
[conda] mkl-service               2.4.0            py39h7f8727e_0  
[conda] mkl_fft                   1.3.1            py39hd3c417c_0  
[conda] mkl_random                1.2.2            py39h51133e4_0  
[conda] numpy                     1.21.2           py39h20f2e39_0  
[conda] numpy-base                1.21.2           py39h79a1101_0  
[conda] pyg                       2.0.3           py39_torch_1.10.0_cu102    pyg
[conda] pytorch                   1.10.2          py3.9_cuda10.2_cudnn7.6.5_0    pytorch
[conda] pytorch-cluster           1.5.9           py39_torch_1.10.0_cu102    pyg
[conda] pytorch-mutex             1.0                        cuda    pytorch
[conda] pytorch-scatter           2.0.9           py39_torch_1.10.0_cu102    pyg
[conda] pytorch-sparse            0.6.12          py39_torch_1.10.0_cu102    pyg
[conda] pytorch-spline-conv       1.2.1           py39_torch_1.10.0_cu102    pyg
[conda] torch-geometric           2.0.4                    pypi_0    pypi
[conda] torchaudio                0.10.2               py39_cu102    pytorch
[conda] torchvision               0.11.3               py39_cu102    pytorch
```

cc @ngimel @jianyuh @nikitaved @pearu @mruberry @walterddr @IvanYashchuk @xwang233 @Lezcano