Update train.py

train.py

@@ -4,7 +4,7 @@ import torch
 from torch import nn 
 from torch.utils.data import DataLoader 
 from torchvision import datasets 
-from torchvision.transforms import ToTensor, Normalize, RandomCrop, RandomHorizontalFlip, Compose 
+from torchvision.transforms import ToTensor, Normalize, Compose 
 from interactor import Interactor
 
 
@@ -43,11 +43,6 @@ for X, y in test_dataloader:
     break
 
 
-def check_sizes(image_size, patch_size):
-    sqrt_num_patches, remainder = divmod(image_size, patch_size)
-    assert remainder == 0, "`image_size` must be divisibe by `patch_size`"
-    num_patches = sqrt_num_patches ** 2
-    return num_patches
 
 
 
@@ -70,7 +65,6 @@ class InteractorImageClassification(Interactor):
        
         
     ):
-        num_patches = check_sizes(image_size, patch_size)
         super().__init__(d_model, num_layers)
         self.patcher = nn.Conv2d(
             in_channels, d_model, kernel_size=patch_size, stride=patch_size