illegal memory access with PointTransformer in torch

I am getting a `RuntimeError: CUDA error: an illegal memory access was encountered` when running the `PointTransformer`. I have tested that on two different systems with nvidia gpu's and got the same error each time.

I am running this in a docker:

```Dockerfile
FROM nvidia/cuda:11.1.1-devel-ubuntu20.04

RUN apt-get update && \
apt-get install -yq --no-install-recommends \
libgl1-mesa-dev \
python3 \
python3-dev \
python3-pip && \
pip3 install --no-cache-dir --upgrade pip && \
rm -rf /var/lib/apt/lists/*

RUN pip install open3d numpy matplotlib tensorboard -U

RUN pip install install https://s3.us-west-1.wasabisys.com/open3d-downloads/torch-1.8.2-cp38-cp38-linux_x86_64.whl \
torchvision==0.9.2+cu111 -f https://download.pytorch.org/whl/lts/1.8/torch_lts.html
```

And the code I am running is:

```python
import open3d.ml as _ml3d
import open3d.ml.torch as ml3d
import numpy as np

data = {
"point": np.random.rand(300000).reshape(-1, 3).astype(np.float32) * 100,
"feat": np.random.rand(300000).reshape(-1, 3).astype(np.float32) * 255,
"label": np.zeros(100000, dtype=np.float32)
}

cfg = _ml3d.utils.Config.load_from_file("pointtransformer_s3dis.yml")
model = ml3d.models.PointTransformer(**cfg.model)
pipeline = ml3d.pipelines.SemanticSegmentation(model, dataset=None, device="cpu", **cfg.pipeline)
pipeline.load_ckpt(ckpt_path="./logs/pointtransformer_s3dis_202109241350utc.pth")
result = pipeline.run_inference(data)
```

Here the full error message:

```
---------------------------------------------------------------------------
RuntimeError Traceback (most recent call last)
/tmp/ipykernel_11/2004900684.py in
3 pipeline = ml3d.pipelines.SemanticSegmentation(model, dataset=None, device="cpu", **cfg.pipeline)
4 pipeline.load_ckpt(ckpt_path="./logs/pointtransformer_s3dis_202109241350utc.pth")
----> 5 result = pipeline.run_inference(data)

open3d/_ml3d/torch/pipelines/semantic_segmentation.py in run_inference(self, data)
165 with torch.no_grad():
166 for unused_step, inputs in enumerate(infer_loader):
--> 167 results = model(inputs['data'])
168 self.update_tests(infer_sampler, inputs, results)
169

torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
887 result = self._slow_forward(*input, **kwargs)
888 else:
--> 889 result = self.forward(*input, **kwargs)
890 for hook in itertools.chain(
891 _global_forward_hooks.values(),

open3d/_ml3d/torch/models/point_transformer.py in forward(self, batch)
173
174 for i in range(5):
--> 175 p, f, r = self.encoders[i]([points[i], feats[i], row_splits[i]])
176 points.append(p)
177 feats.append(f)

torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
887 result = self._slow_forward(*input, **kwargs)
888 else:
--> 889 result = self.forward(*input, **kwargs)
890 for hook in itertools.chain(
891 _global_forward_hooks.values(),

torch/nn/modules/container.py in forward(self, input)
117 def forward(self, input):
118 for module in self:
--> 119 input = module(input)
120 return input
121

torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
887 result = self._slow_forward(*input, **kwargs)
888 else:
--> 889 result = self.forward(*input, **kwargs)
890 for hook in itertools.chain(
891 _global_forward_hooks.values(),

open3d/_ml3d/torch/models/point_transformer.py in forward(self, pxo)
642 identity = feat
643 feat = self.relu(self.bn1(self.linear1(feat)))
--> 644 feat = self.relu(self.bn2(self.transformer2([point, feat, row_splits])))
645 feat = self.bn3(self.linear3(feat))
646 feat += identity

torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
887 result = self._slow_forward(*input, **kwargs)
888 else:
--> 889 result = self.forward(*input, **kwargs)
890 for hook in itertools.chain(
891 _global_forward_hooks.values(),

open3d/_ml3d/torch/models/point_transformer.py in forward(self, pxo)
429 feat_q, feat_k, feat_v = self.linear_q(feat), self.linear_k(
430 feat), self.linear_v(feat) # (n, c)
--> 431 feat_k = queryandgroup(self.nsample,
432 point,
433 point,

open3d/_ml3d/torch/models/point_transformer.py in queryandgroup(nsample, points, queries, feat, idx, points_row_splits, queries_row_splits, use_xyz)
679 queries = points
680 if idx is None:
--> 681 idx = knn_batch(points,
682 queries,
683 k=nsample,

open3d/_ml3d/torch/models/point_transformer.py in knn_batch(points, queries, k, points_row_splits, queries_row_splits, return_distances)
733 -1, k).long().cuda(), ans.neighbors_distance.reshape(-1, k).cuda()
734 else:
--> 735 return ans.neighbors_index.reshape(-1, k).long().cuda()
736
737

RuntimeError: CUDA error: an illegal memory access was encountered
```