KPConv on tensorflow produces large loss when using deformable layers

### Checklist

- [X] I have searched for [similar issues](https://github.com/isl-org/Open3D-ML/issues).
- [X] I have tested with the [latest development wheel](http://www.open3d.org/docs/latest/getting_started.html#development-version-pip).
- [X] I have checked the [release documentation](http://www.open3d.org/docs/release/) and the [latest documentation](http://www.open3d.org/docs/latest/) (for `master` branch).

### Describe the issue

I'm trying to train KPConv on the Paris-Lille3D dataset but I'm getting a very high loss (in the millions range) right from the start, which makes training impossible - after ~100 epochs there is no negative trend. I am using the default config `kpconv_parislille3d.yml`. After switching the architecture to not use any deformable layers training works as expected, which leads me to believe there is a problem with the implementation of the deformable layer.

<!--table
{mso-displayed-decimal-separator:"\.";
mso-displayed-thousand-separator:"\,";}
@page
{margin:.75in .7in .75in .7in;
mso-header-margin:.3in;
mso-footer-margin:.3in;}
tr
{mso-height-source:auto;}
col
{mso-width-source:auto;}
br
{mso-data-placement:same-cell;}
td
{padding-top:1px;
padding-right:1px;
padding-left:1px;
mso-ignore:padding;
color:black;
font-size:11.0pt;
font-weight:400;
font-style:normal;
text-decoration:none;
font-family:Calibri, sans-serif;
mso-font-charset:0;
mso-number-format:General;
text-align:general;
vertical-align:bottom;
border:none;
mso-background-source:auto;
mso-pattern:auto;
mso-protection:locked visible;
white-space:nowrap;
mso-rotate:0;}
.xl65
{text-align:center;
vertical-align:middle;
white-space:normal;}
-->

architecture, which works | problematic architecture
-- | --
simple | simple
resnetb | resnetb
resnetb_strided | resnetb_strided
resnetb | resnetb
resnetb | resnetb_strided
resnetb_strided | resnetb_deformable
resnetb | resnetb_deformable_strided
resnetb | resnetb_deformable
resnetb_strided | resnetb_deformable_strided
resnetb | resnetb_deformable
resnetb | nearest_upsample
resnetb_strided | unary
resnetb | nearest_upsample
nearest_upsample | unary
unary | nearest_upsample
nearest_upsample | unary
unary | nearest_upsample
nearest_upsample | unary
unary |  
nearest_upsample |  
unary |  

### Steps to reproduce the bug

```python
1. Download the dataset from https://npm3d.fr/paris-lille-3d
2. On Colab run the following cell:

!pip uninstall -y -qqq tensorflow
!pip install -qqq tensorflow==2.5.3
!pip install -qqq open3d
!pip uninstall -y -qqq tensorflow-probability
!pip install -qqq tensorflow-probability==0.13.0
!pip uninstall -y numpy
!pip install numpy==1.19.5
```
3. Finally run the training:

```
import os
import open3d.ml as _ml3d
import open3d.ml.tf as ml3d

cfg_file = config_path
cfg = _ml3d.utils.Config.load_from_file(cfg_file)

model = ml3d.models.KPFCNN(**cfg.model)
cfg.dataset['dataset_path'] = dataset_path
dataset = ml3d.datasets.ParisLille3D(cfg.dataset.pop('dataset_path', None), **cfg.dataset)
pipeline = ml3d.pipelines.SemanticSegmentation(model, dataset=dataset, device="gpu", **cfg.pipeline)

# prints training progress in the console.
pipeline.run_train()
```
```

### Error message

INFO - 2022-05-17 17:06:02,791 - semantic_segmentation -
INFO - 2022-05-17 17:06:02,792 - semantic_segmentation - Logging in file : ./logs/KPFCNN_ParisLille3D_tf/log_train_2022-05-17_17:06:02.txt
INFO - 2022-05-17 17:06:02,801 - parislille3d - Found 3 pointclouds for training
INFO - 2022-05-17 17:06:05,246 - parislille3d - Found 1 pointclouds for validation
INFO - 2022-05-17 17:06:07,220 - semantic_segmentation - Writing summary in train_log/00010_KPFCNN_ParisLille3D_tf.
INFO - 2022-05-17 17:06:07,231 - semantic_segmentation - Initializing from scratch.
INFO - 2022-05-17 17:06:07,235 - semantic_segmentation - === EPOCH 0/100 ===
training: 41it [01:51, 2.73s/it]
validation: 11it [00:13, 1.19s/it]
INFO - 2022-05-17 17:08:12,172 - semantic_segmentation - loss train: 6508379.000 eval: 528964255165186048.000
INFO - 2022-05-17 17:08:12,177 - semantic_segmentation - acc train: 0.107 eval: 0.105
INFO - 2022-05-17 17:08:12,182 - semantic_segmentation - iou train: 0.052 eval: 0.042
INFO - 2022-05-17 17:08:13,295 - semantic_segmentation - Saved checkpoint at: ./logs/KPFCNN_ParisLille3D_tf/checkpoint/ckpt-1
INFO - 2022-05-17 17:08:13,297 - semantic_segmentation - === EPOCH 1/100 ===
training: 100%|██████████| 40/40 [01:55<00:00, 1.58s/it]

### Expected behavior

The expected behavior would be a loss in the range 0