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

start = time()

num_epochs = 5
for epoch in range(1, num_epochs + 1):
    train(model, device, train_loader, optimizer, criterion, epoch)

print(f"Total time = {time() - start}")