How do I implement the following pytorch operations in burn?

[zemelLeong]

>>> tensor_randn.shape
torch.Size([1, 54, 8])
>>> tensor_randn[:, :, 0].shape
torch.Size([1, 54])

[nathanielsimard]

let tensor = Tensor::random(..);
let [b, d1, d2] = tensor.dims();
let tensor_partial = tensor.slice([0..b, 0..d1, 0..1]);

You have to specify the range of each dimension; at some point, we might have syntax sugar to avoid specifying the start and end position of dimensions where you want to keep everything, maybe something like that:

let tensor_partial = tensor.slice(index![.., .., 0]);

We might also implement the Index trait with multiple different generics: https://doc.rust-lang.org/std/ops/trait.Index.html, dispatching to slice when passing Range and to select when using an int tensor.

[zemelLeong]

Thank you. A temporary solution is written here for those in need to refer to.

let tensor = Tensor::random(..);
let [b, d1, d2] = tensor.dims();
let tensor_partial = tensor.slice([0..b, 0..d1, 0..1]).reshape([b, -1]);
// tensor_partial.shape [b, d1]