Consegui resolver o problema.
Na configuração de hiperparâmetros basta mudar o num_workers para 0 e irá rodar normalmente.
Ficará da seguinte forma:
# Configurando hiperparâmetros.
args = {
'epoch_num': 200, # Número de épocas.
'lr': 5e-5, # Taxa de aprendizado.
'weight_decay': 5e-4, # Penalidade L2 (Regularização).
'num_workers': 0, # Número de threads do dataloader.
'batch_size': 20, # Tamanho do batch.
}
Encontrei nos fóruns que é um problema relacionado ao ambiente Windows e que essa mudança não gera uma perca significativa na eficiencia do código, mas ninguém explicou qual é o problema propriamente dito.
E faz sentido o código rodar no Colab com num_workers = 3, pois o problema é com o Windows.