I misunderstood this part:



 Luca_Pamparana:

However, the dimension of the input tensor is dynamic


as x would be the input tensor.
Note that rng is a 1-dimensional tensor with 3 values.
If you want to index the last 3 values, you could use negative indices:
x[:, :, :rng[-3], :rng[-2], :rng[-1]]