Input key error in 3d object detection example

Using master (12/04/2021) + bypassing tensor_requirements to be 2.4.1+ on ubuntu 20.04 = successful data visualization on data reader example using semantic KITTI dataset.

After downloading the KITTI data/label/calib as per dataset python script, I'm running the following example for the point pillars

> import os
> import open3d.ml as _ml3d

> import open3d.ml.tf as ml3d
>
> cfg_file = "/content/point_pillars/Open3D-ML/ml3d/configs/pointpillars_kitti.yml"
> cfg = _ml3d.utils.Config.load_from_file(cfg_file)
>
> model = ml3d.models.PointPillars(**cfg.model)
> cfg.dataset['dataset_path'] = "/content/KITTI_DATASET/KITTI_PTCLOUD_DATA/data_object_velodyne"
> dataset = ml3d.datasets.KITTI(cfg.dataset.pop('dataset_path', None), **cfg.dataset)
>
> pipeline = ml3d.pipelines.ObjectDetection(model, dataset=dataset, device="gpu", **cfg.pipeline)
>
> ckpt_path = '/content/KITTI_DATASET/pointpillars_kitti_202012221652utc/ckpt-12'

> pipeline.load_ckpt(ckpt_path=ckpt_path)
>
> test_split = dataset.get_split("test")
>
> data = test_split.get_data(0)

> result = pipeline.run_inference(data)
>
> pipeline.run_test()
> '''

Yields the following error (most of them warnings for not having CUDA):

> 2021-12-04 21:22:00.545338: W tensorflow/stream_executor/platform/default/dso_loader.cc:60] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory
> 2021-12-04 21:22:00.545360: I tensorflow/stream_executor/cuda/cudart_stub.cc:29] Ignore above cudart dlerror if you do not have a GPU set up on your machine.
> 2021-12-04 21:22:01.339405: W tensorflow/stream_executor/platform/default/dso_loader.cc:60] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory
> 2021-12-04 21:22:01.339424: I tensorflow/stream_executor/cuda/cudart_stub.cc:29] Ignore above cudart dlerror if you do not have a GPU set up on your machine.
> 2021-12-04 21:22:01.908347: I tensorflow/compiler/jit/xla_cpu_device.cc:41] Not creating XLA devices, tf_xla_enable_xla_devices not set
> 2021-12-04 21:22:01.908474: W tensorflow/stream_executor/platform/default/dso_loader.cc:60] Could not load dynamic library 'libcuda.so.1'; dlerror: libcuda.so.1: cannot open shared object file: No such file or directory
> 2021-12-04 21:22:01.908485: W tensorflow/stream_executor/cuda/cuda_driver.cc:326] failed call to cuInit: UNKNOWN ERROR (303)
> 2021-12-04 21:22:01.908501: I tensorflow/stream_executor/cuda/cuda_diagnostics.cc:156] kernel driver does not appear to be running on this host (dsohn03): /proc/driver/nvidia/version does not exist
> 2021-12-04 21:22:01.908649: I tensorflow/core/platform/cpu_feature_guard.cc:142] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations: AVX2 FMA
> To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.
> 2021-12-04 21:22:01.909075: I tensorflow/compiler/jit/xla_gpu_device.cc:99] Not creating XLA devices, tf_xla_enable_xla_devices not set
> INFO - 2021-12-04 21:22:02,104 - object_detection - Restored from /home/daniel.sohn/Desktop/KITTI_DATASET/pointpillars_kitti_202012221652utc/ckpt-12
> INFO - 2021-12-04 21:22:02,105 - kitti - Found 7518 pointclouds for test
> Traceback (most recent call last):
> File "./demo_ptpillars.py", line 37, in
> result = pipeline.run_inference(data)
> File "/content/point_pillars/env/lib/python3.8/site-packages/open3d/_ml3d/tf/pipelines/object_detection.py", line 57, in run_inference
> results = model(data, training=False)
> File "/content/point_pillars/env/lib/python3.8/site-packages/tensorflow/python/keras/engine/base_layer.py", line 1012, in __call__
> outputs = call_fn(inputs, *args, **kwargs)
> File "/content/point_pillars/env/lib/python3.8/site-packages/open3d/_ml3d/tf/models/point_pillars.py", line 142, in call
> inputs = unpack(inputs[0], inputs[-2])
> KeyError: 0

Is this a repeat issue of https://github.com/isl-org/Open3D-ML/issues/408 ?

tensorflow 2.4.4
open3d 0.13.0

Thank you,