evaluating at different time points per batch

[LucaCras]

Hi,

I've got a tensor of the shape batch_size x T x nr_channels and I created a cubic spline accordingly.

Let's say I have created this tensor with shape (2, 10, 1)

Now I want to query [[0], [1]] (e.g. value at t=0 for batch 0 and value at t=1 for batch 1)

I know I can query spline.evaluate(torch.tensor(0)) to get the value at t=0 for both batches, but how can I query the above such that it returns a tensor of shape (2, 1, 1) or (2, 1) as the above query returns (2, 2, 1) or even (2, 2, 1, 1).