What were the print statements returning in the iteration before this error is raised?
Did you make sure that edge_index[idx] returned a valid tensor? If so, could you post an executable code snippet to reproduce this issue?