train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True, num_workers=8)

start = time()

num_epochs = 5

for epoch in range(1, num_epochs + 1):
    train(model, device, train_loader, optimizer, criterion, epoch, non_blocking=True)
print(f"Total time = {time() - start}")