You can do a outer for loop and just look at input = full_input.select(0, batch_idx) every time.
Note that select() does not copy memory and so any inplace change of input will be reflected into full_input !