#手っ取り早い解決方法
以下のようにdrop_lastをTrueで設定してあげましょう。
train_loader = torch.utils.data.DataLoader(
train_dataset,
batch_size=batch_size,
shuffle=True,
num_workers=workers,
pin_memory=True,
drop_last=True,
sampler=None)
#<参考>
このエラーの原因はラストのミニバッチの画像クラスがどちらか一方しかないことです。(下により詳しく載ってます)
ROC,AUCについてはここに詳しく載ってます。