When you are calling state = all_state_torch[st:st+batch] it is simple operation of creating view over all_state_torch tensor. Which is super fast in compare to DataLoader because it creates completely new tensor as concatenation of multiple self.state[index] pieces.