def train(data_loader):
    start = time.time()
    for _ in tqdm(range(10)):
        for x in data_loader:
            pass
    end = time.time()
    return end - start

This part is super confusing. Is this intentional that you want to iterate over DataLoader 10 times?