Is it possible to export stgcn_wave to a torchscript model?

## 🐛 Bug

## To Reproduce

I want to export [stgcn_wave](https://github.com/dmlc/dgl/tree/master/examples/pytorch/stgcn_wave) model to a torchscript model.
What I did is:

1. add these two lines below [the line 11](https://github.com/dmlc/dgl/blob/master/examples/pytorch/stgcn_wave/utils.py#L11)
```
traced_script_module = torch.jit.trace(model, x, check_trace=False)
traced_script_module.save("stgcn_torchscript.pt")
```
2. then I got this error:
```
/lib/python3.6/site-packages/dgl/nn/pytorch/conv/graphconv.py:388: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
if (graph.in_degrees() == 0).any():
/lib/python3.6/site-packages/dgl/heterograph.py:2730: TracerWarning: Using len to get tensor shape might cause the trace to be incorrect. Recommended usage would be tensor.shape[0]. Passing a tensor of different shape might lead to errors or silently give incorrect results.
if len(vid_tensor) > 0 and F.as_scalar(F.min(vid_tensor, 0)) < 0 < len(vid_tensor):
/lib/python3.6/site-packages/dgl/backend/pytorch/tensor.py:49: TracerWarning: Converting a tensor to a Python number might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
return data.item()
/lib/python3.6/site-packages/dgl/heterograph.py:3580: TracerWarning: Using len to get tensor shape might cause the trace to be incorrect. Recommended usage would be tensor.shape[0]. Passing a tensor of different shape might lead to errors or silently give incorrect results.
if F.as_scalar(F.sum(self.has_nodes(u_tensor, ntype=srctype), dim=0)) != len(u_tensor):
/lib/python3.6/site-packages/dgl/heterograph.py:4109: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
if nfeats != num_nodes:
Traceback (most recent call last):
File "main.py", line 117, in
val_loss = evaluate_model(model, loss, val_iter)
File "/utils.py", line 14, in evaluate_model
traced_script_module.save("stgcn_torchscript.pt")
File "/lib/python3.6/site-packages/torch/jit/_script.py", line 596, in save
return self._c.save(str(f), **kwargs)
RuntimeError:
Could not export Python function call 'GSpMM'. Remove calls to Python functions before export. Did you forget to add @script or @script_method annotation? If this is a nn.ModuleList, add it to __constants__:
/lib/python3.6/site-packages/dgl/backend/pytorch/sparse.py(502): gspmm
/lib/python3.6/site-packages/dgl/ops/spmm.py(77): gspmm
/lib/python3.6/site-packages/dgl/ops/spmm.py(191): func
/lib/python3.6/site-packages/dgl/core.py(297): invoke_gspmm
/lib/python3.6/site-packages/dgl/core.py(322): message_passing
/lib/python3.6/site-packages/dgl/heterograph.py(4849): update_all
/lib/python3.6/site-packages/dgl/nn/pytorch/conv/graphconv.py(434): forward
/lib/python3.6/site-packages/torch/nn/modules/module.py(1039): _slow_forward
/lib/python3.6/site-packages/torch/nn/modules/module.py(1051): _call_impl
/model.py(46): forward
/lib/python3.6/site-packages/torch/nn/modules/module.py(1039): _slow_forward
/lib/python3.6/site-packages/torch/nn/modules/module.py(1051): _call_impl
/model.py(100): forward
/lib/python3.6/site-packages/torch/nn/modules/module.py(1039): _slow_forward
/lib/python3.6/site-packages/torch/nn/modules/module.py(1051): _call_impl
/lib/python3.6/site-packages/torch/jit/_trace.py(959): trace_module
/lib/python3.6/site-packages/torch/jit/_trace.py(744): trace
/utils.py(13): evaluate_model
main.py(117):

```

## Expected behavior

I want to save the model into a torchscript model, then I will run this torchscript model in java. [See this tutorial.](https://docs.djl.ai/jupyter/load_pytorch_model.html#step-1-prepare-your-model)

## Environment

- dgl-cu102: 0.7.2
- torch: 1.9.0
- OS: ubuntu 16.04
- How you installed DGL (`conda`, `pip`, source): pip
- Python version: 3.6
- CUDA/cuDNN version (if applicable): 10.2
- GPU models and configuration: GeForce GTX 1080