Hello,
In my shallow opinion, there might be something wrong with your snippet.



 zelazny:

idx = train_dataset.targets==i


Here, idx is tensor with values 0 and 1. So if you pass it as indices argument to Subset, the subset will only contain the first and the second target (means position 0 and position 1).
And train_dataset.targets[idx] does not make sense to me.



 zelazny:

only one element tensors can be converted to Python scalars


For this error, you could try to print train_iters[lable_id] 's shape and debug it.