Thank you very much for your comment.
The reason im feeding the samples one by one from the input tensor, is that based on the extra_label, i want to use different branch for forward the input.
I understand now why the aproach is incorrect.
I will handle somehow that i only feed a tensor of the same extra_label, in order to only feed it the whole tensor to one branch.
I appreciate the time you spent in order to help me! Thank you sir!!