Corrected the torch.transforms

data_setup.py

@@ -7,17 +7,17 @@ import numpy as np
 import json
 import torch.nn as nn
 import torch.nn.functional as F
-# import torchvision.transforms as tt
-import albumentations as A
-from albumentations.pytorch import ToTensorV2
+import torchvision.transforms as T
+# import albumentations as A
+# from albumentations.pytorch import ToTensorV2
 
 
 stats = (0.4862, 0.4561, 0.3941), (0.2202, 0.2142, 0.2160)
 
-model_tsfm = A.Compose([
-                A.Resize(224, 224),
+model_tsfm = T.Compose([
+                T.Resize((224, 224)),
                 # A.Normalize(*stats),
-                ToTensorV2()             
+                T.ToTensor()             
              ])
 
 
@@ -29,7 +29,6 @@ cat_to_name.index = cat_to_name.index.astype(np.int32)
 cat_to_name.sort_index(inplace=True)
 classes = cat_to_name.values
 
-# classes = ['Australian terrier', 'Border terrier', 'Samoyed', 'Beagle', 'Shih-Tzu', 'English foxhound', 'Rhodesian ridgeback', 'Dingo', 'Golden retriever', 'Old English sheepdog']
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()