Confused: A vital error about TypedLinear function and Single machine model parallel best practices

## 🐛 Bug

## To Reproduce
I'm doing Single machine model parallel best practices with a tutorial from PyTorch and its Basic Usage [https://pytorch.org/tutorials/intermediate/model_parallel_tutorial.html](url)

And I found the function "dgl.nn.pytorch.TypedLinear" is incompatible, and will raise an exception error.

Here is a simple example code I created:

```
from dgl.nn import TypedLinear
import torch
import torch.nn as nn
import torch.optim as optim

class ToyModel(nn.Module):
def __init__(self):
super(ToyModel, self).__init__()
self.net1 = TypedLinear(32, 64, 5).to('cuda:0')
self.relu = torch.nn.ReLU()
self.net2 = TypedLinear(64, 64, 5).to('cuda:1')

def forward(self, x, x_type):
x = self.relu(self.net1(x.to('cuda:0'), x_type.to('cuda:0')))
return self.net2(x.to('cuda:1'), x_type.to('cuda:1'))

x = torch.randn(100, 32)
x_type = torch.randint(0, 5, (100,))

model = ToyModel()
for i in range(10): #here is a loop just like when we train a model
print(i)
y = model(x, x_type)
print(y.shape)
```
And then I will get an error:
```
---------------------------------------------------------------------------
RuntimeError Traceback (most recent call last)
Cell In[2], line 14
12 for i in range(10):
13 print(i)
---> 14 y = model(x, x_type)
15 print(y.shape)

File ~/anaconda3/envs/lyk/lib/python3.9/site-packages/torch/nn/modules/module.py:1130, in Module._call_impl(self, *input, **kwargs)
1126 # If we don't have any hooks, we want to skip the rest of the logic in
1127 # this function, and just call forward.
1128 if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
1129 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1130 return forward_call(*input, **kwargs)
1131 # Do not call functions when jit is used
1132 full_backward_hooks, non_full_backward_hooks = [], []

Cell In[2], line 8, in ToyModel.forward(self, x, x_type)
7 def forward(self, x, x_type):
----> 8 x = self.relu(self.net1(x.to('cuda:0'), x_type.to('cuda:0')))
9 return self.net2(x.to('cuda:1'), x_type.to('cuda:1'))

RuntimeError: CUDA error: an illegal memory access was encountered
CUDA kernel errors might be asynchronously reported at some other API call,so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.
```

if I add

```
import os
os.environ['CUDA_LAUNCH_BLOCKING'] = '1'
```
in the beginning of the code. The error will be like
```
---------------------------------------------------------------------------
DGLError Traceback (most recent call last)
Cell In[2], line 16
14 for i in range(10):
15 print(i)
---> 16 y = model(x, x_type)
17 print(y.shape)

File ~/anaconda3/envs/lyk/lib/python3.9/site-packages/torch/nn/modules/module.py:1130, in Module._call_impl(self, *input, **kwargs)
1126 # If we don't have any hooks, we want to skip the rest of the logic in
1127 # this function, and just call forward.
1128 if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
1129 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1130 return forward_call(*input, **kwargs)
1131 # Do not call functions when jit is used
1132 full_backward_hooks, non_full_backward_hooks = [], []

Cell In[2], line 9, in ToyModel.forward(self, x, x_type)
7 def forward(self, x, x_type):
8 x = self.relu(self.net1(x, x_type))
----> 9 return self.net2(x.to('cuda:1'), x_type.to('cuda:1'))

File ~/anaconda3/envs/lyk/lib/python3.9/site-packages/torch/nn/modules/module.py:1130, in Module._call_impl(self, *input, **kwargs)
1126 # If we don't have any hooks, we want to skip the rest of the logic in
1127 # this function, and just call forward.
1128 if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
1129 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1130 return forward_call(*input, **kwargs)
1131 # Do not call functions when jit is used
1132 full_backward_hooks, non_full_backward_hooks = [], []

File ~/anaconda3/envs/lyk/lib/python3.9/site-packages/dgl/nn/pytorch/linear.py:174, in TypedLinear.forward(self, x, x_type, sorted_by_type)
172 return segment_mm(x, w, seglen_a=seglen)
173 else:
--> 174 return gather_mm(x, w, idx_b=x_type)

File ~/anaconda3/envs/lyk/lib/python3.9/site-packages/dgl/ops/gather_mm.py:40, in gather_mm(a, b, idx_b)
38 pos_r = torch.cat([pos_l[1:], torch.tensor([len(idx_b)], device=a.device)])
39 seglen = (pos_r - pos_l).cpu() # XXX(minjie): cause device synchronize
---> 40 return torch.index_select(F.segment_mm(sorted_a, b, seglen), 0, rev_perm)
41 else:
42 return F.gather_mm(a, b, None, idx_b)

File ~/anaconda3/envs/lyk/lib/python3.9/site-packages/dgl/backend/pytorch/sparse.py:802, in segment_mm(A, B, seglen_A)
800 return th.cat(C)
801 else:
--> 802 return SEGMENTMM.apply(A, B, seglen_A)

File ~/anaconda3/envs/lyk/lib/python3.9/site-packages/torch/cuda/amp/autocast_mode.py:118, in custom_fwd..decorate_fwd(*args, **kwargs)
116 return fwd(*_cast(args, cast_inputs), **_cast(kwargs, cast_inputs))
117 else:
--> 118 return fwd(*args, **kwargs)

File ~/anaconda3/envs/lyk/lib/python3.9/site-packages/dgl/backend/pytorch/sparse.py:671, in SEGMENTMM.forward(ctx, A, B, seglen_A)
669 raise ValueError("segment_mm expects B to be a 3D tensor.")
670 C = th.zeros((A.shape[0], B.shape[2]), device=A.device, dtype=A.dtype)
--> 671 C = _segment_mm(A, B, C, seglen_A)
672 ctx.backward_cache = A, B, seglen_A
673 return C

File ~/anaconda3/envs/lyk/lib/python3.9/site-packages/dgl/sparse.py:394, in _segment_mm(A, B, out, seglen_A, b_trans)
392 def _segment_mm(A, B, out, seglen_A, b_trans=False):
393 """Invoke the C API of segment_mm."""
--> 394 _CAPI_DGLKernelSEGMENTMM(to_dgl_nd(A),
395 to_dgl_nd(B),
396 to_dgl_nd_for_write(out),
397 to_dgl_nd(seglen_A),
398 False, b_trans)
399 return out

File dgl/_ffi/_cython/./function.pxi:293, in dgl._ffi._cy3.core.FunctionBase.__call__()

File dgl/_ffi/_cython/./function.pxi:239, in dgl._ffi._cy3.core.FuncCall()

DGLError: [19:57:01] /opt/dgl/src/array/cuda/gather_mm.cu:247: Check failed: e == CUBLAS_STATUS_SUCCESS: CUBLAS ERROR: 13
Stack trace:
[bt] (0) /home/cseadmin/lyk/anaconda3/envs/lyk/lib/python3.9/site-packages/dgl/libdgl.so(dmlc::LogMessageFatal::~LogMessageFatal()+0x4f) [0x7f0ebe22ad6f]
[bt] (1) /home/cseadmin/lyk/anaconda3/envs/lyk/lib/python3.9/site-packages/dgl/libdgl.so(void dgl::aten::SegmentMM<2, long, 32>(dgl::runtime::NDArray, dgl::runtime::NDArray, dgl::runtime::NDArray, dgl::runtime::NDArray, bool, bool)+0x393) [0x7f0ebe762bb3]
[bt] (2) /home/cseadmin/lyk/anaconda3/envs/lyk/lib/python3.9/site-packages/dgl/libdgl.so(dgl::aten::SegmentMM(dgl::runtime::NDArray, dgl::runtime::NDArray, dgl::runtime::NDArray, dgl::runtime::NDArray, bool, bool)+0x14e7) [0x7f0ebe4dfe77]
[bt] (3) /home/cseadmin/lyk/anaconda3/envs/lyk/lib/python3.9/site-packages/dgl/libdgl.so(+0x663a39) [0x7f0ebe4faa39]
[bt] (4) /home/cseadmin/lyk/anaconda3/envs/lyk/lib/python3.9/site-packages/dgl/libdgl.so(DGLFuncCall+0x48) [0x7f0ebe555928]
[bt] (5) /home/cseadmin/lyk/anaconda3/envs/lyk/lib/python3.9/site-packages/dgl/_ffi/_cy3/core.cpython-39-x86_64-linux-gnu.so(+0x16cb0) [0x7f0ebddfccb0]
[bt] (6) /home/cseadmin/lyk/anaconda3/envs/lyk/lib/python3.9/site-packages/dgl/_ffi/_cy3/core.cpython-39-x86_64-linux-gnu.so(+0x1741b) [0x7f0ebddfd41b]
[bt] (7) /home/cseadmin/lyk/anaconda3/envs/lyk/bin/python(_PyObject_MakeTpCall+0x2df) [0x557ed552547f]
[bt] (8) /home/cseadmin/lyk/anaconda3/envs/lyk/bin/python(_PyEval_EvalFrameDefault+0x417) [0x557ed55bed57]
```

## Expected behavior
If correct, the output should be correct each time through the loop

## Environment

- DGL Version (e.g., 1.0): **every Version>0.8.1**
- Backend Library & Version (e.g., PyTorch 0.4.1, MXNet/Gluon 1.3): **PyTorch>1.9.1**
- OS (e.g., Linux):
- How you installed DGL (`conda`, `pip`, source): both conda and Pip
- Build command you used (if compiling from source):
- Python version: **3.7.7**
- CUDA/cuDNN version (if applicable):
- GPU models and configuration (e.g. V100): **2080Ti**
- Any other relevant information:

## Additional context
**I am really confused as to why this is happening and hope it can be solved or fixed as soon as possible**

Reproduce the two-GPU example from the issue, then inspect dgl/nn/pytorch/linear.py and the gather_mm, segment_mm, and _segment_mm paths named in the traceback. Compare behavior across the reported DGL and PyTorch versions; done means TypedLinear runs repeatedly in the single-machine model-parallel example without CUDA or CUBLAS errors.

Written by the indexing model from the issue text.