First, you can set batch_size for the DataLoader to load multiple data into a batch and do transform or augmentation for each batch.
Secondly, I am not sure why you have this tmp_list. You can directly use the augmented data to train your model.
for x, y in train_loader:
    aug_x = aug(x)
    ...
    output = model(aug_x)
    loss = ...
    loss.backward()
    ...
``