General Union types are not currently supported in v1.6.1

[Albert zhao]

RuntimeError:
General Union types are not currently supported. Only Union[T, NoneType] (i.e. Optional[T]) is supported.:
File "/home/zj/miniconda3/envs/yukewang/lib/python3.8/site-packages/torch_cluster/rw.py", line 18
num_nodes: Optional[int] = None,
return_edge_indices: bool = False,
) -> Union[Tensor, Tuple[Tensor, Tensor]]:
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
"""Samples random walks of length :obj:walk_length from all node indices
in :obj:start in the graph given by :obj:(row, col) as described in the

[Matthias Fey]

What PyTorch version are you on?