Error when using @tf.function in release 0.17.0

[siebertdarius]

Hello Nvidia,

I have an issue using the @tf.function in the new sionna release. The error (see below) occured and can be reproduced when running the notebook code for learned materials from your diff-rt-calibration project with sionna 0.17.0 release.

This code works perfectly for sionna 0.16.2. Also, the error refers to the doppler shifts which are touched in release 0.17.0, so I suspect an issue within this new release.

When @tf.function is deleted in the code, no error occurs but the results differ drastically from the usual results.

    TypeError: in user code:

    File "/tmp/ipykernel_481/3933607846.py", line 61, in training_step  *
        paths = scene.compute_fields(*traced_paths,
    File "/opt/conda/lib/python3.11/site-packages/sionna/rt/scene.py", line 620, in compute_fields  *
        output = self._solver_paths.compute_fields(spec_paths, diff_paths,
    File "/opt/conda/lib/python3.11/site-packages/sionna/rt/solver_paths.py", line 754, in compute_fields  *
        spec_paths.doppler = self._compute_doppler_shifts(spec_paths,
    File "/opt/conda/lib/python3.11/site-packages/sionna/rt/solver_paths.py", line 4127, in _compute_doppler_shifts  *
        doppler = tf.where(objects_mask, 0, doppler)

    TypeError: Input 'e' of 'SelectV2' Op has type float32 that does not match type int32 of argument 't'.

[SebastianCa]

Hi @siebertdarius,

thank you for reporting this issue. The issue will be fixed in the next release of Sionna (see #436).