What is the Pytorch sub for this tensor flow code?
In converting this line of code to Pytorch from Tensor Flow, I am having trouble
datagen =  ImageDataGenerator(
  shear_range=0.2,
  zoom_range=0.2,
)
def read_img(filename, size, path):
    img = image.load_img(os.path.join(path, filename), target_size=size)
    #convert image to array
    img = img_to_array(img) / 255
    return img

and then
corona_df = final_train_data[final_train_data['Label_2_Virus_category'] == 'COVID-19']
with_corona_augmented = []

#create a function for augmentation
def augment(name):
    img = read_img(name, (255,255), train_img_dir)
    i = 0
    for batch in tqdm(datagen.flow(tf.expand_dims(img, 0), batch_size=32)):
        with_corona_augmented.append(tf.squeeze(batch).numpy())
        if i == 20:
            break
        i =i+1

#apply the function
corona_df['X_ray_image_name'].apply(augment)

I tried doing
transform = transforms.Compose([transforms.Resize(255*255)
                               ])
train_loader = torch.utils.data.DataLoader(os.path.join(train_dir,corona_df),transform = transform,batch_size =32)
def read_img(path):
    img  = train_loader()
    img = np.asarray(img,dtype='int32')
    img = img/255
    return img

I tried continuing but got soo confused by the errors.
I welcome any feedback. Tell me If i miss something
Even a small advice would work, thanks !

You can create a custom dataset to read the images. If you have a directory full of images you can use ImageFolder default dataset. Otherwise if you have different folder placement you can write your own custom dataset class. You can look to this link for custom datasets. What dataloader does is, it automatically gets the data from your dataset and read the images according to your dataset __getitem__ function and apply transformation. So you don't need anything fancy to apply augmentation.
transform = transforms.Compose([ transforms.RandomAffine(20,shear=20,scale=(-0.2,0.2)),
                                 transforms.Resize(255*255)
                               ])

dataset = torchvision.datasets.ImageFolder(train_img_dir, transform=transform)
loader = torch.utils.data.DataLoader(dataset,batch_size =32,shuffle=True)

for batch in loader:
    output = model(batch)

