I don’t see the error in the linked notebook, but would guess that names[i] is failing for i=1.
Print names right before calling the plot function and make sure it’s containing the expected values.
Your second issue sounds as if the model is overfitting to a specific class. I’m not familiar with your use case but in case you are working with an imbalanced dataset, try to use a WeightedRandomSampler if possible or a weighted loss. It’s also a good idea to try to overfit a small dataset first (e.g. just 10 samples) by playing around with the hyperparameters to see if any obvious bugs are still in the code.