Is there a way for me to select these rows and form a tensor?
yes, you can use PyTorch indexing to select specific rows.
For example:
output = x[[0, 1, 2], [1, 2, 3], :]
