Update model.py

Previous
import torch
import torchvision

from torch import nn


def create_ViT(num_classes:int=126,
seed:int=42):

weights = torchvision.models.ViT_B_16_Weights.DEFAULT
transforms = weights.transforms()
model = torchvision.models.vit_b_16(weights=weights)

for param in model.parameters():
param.requires_grad = False

torch.manual_seed(seed)
model.heads = nn.Sequential(nn.Linear(in_features=768, out_features=num_classes))

return model, transforms

model.py

@@ -1,20 +1,8 @@
-import torch
-import torchvision
+from transformers import ViTFeatureExtractor, ViTForImageClassification
 
-from torch import nn
+def create_ViT():
 
+  extractor = ViTFeatureExtractor.from_pretrained("google/vit-base-patch16-224")
+  model = ViTForImageClassification.from_pretrained("google/vit-base-patch16-224")
 
-def create_ViT(num_classes:int=126,
-                    seed:int=42):
-
-  weights = torchvision.models.ViT_B_16_Weights.DEFAULT
-  transforms = weights.transforms()
-  model = torchvision.models.vit_b_16(weights=weights)
-
-  for param in model.parameters():
-      param.requires_grad = False
-
-  torch.manual_seed(seed)
-  model.heads = nn.Sequential(nn.Linear(in_features=768, out_features=num_classes))
-
-  return model, transforms
+  return model