Thanks for the reply. I realized after seeing your code that I had one extra dimension in my weights which was why my dataset __getitem__was getting a list of indices instead of single index at a time.
Had to change
weights[i] = weight_per_class[np.where(unique_labels == val)[0]]
to
weights[i] = weight_per_class[np.where(unique_labels == val)[0][0]]
in my stratified sampling function and it works now.
Thank you.