import torch, torchvision def create_vitB16_model(num_classes: int=3, seeds: int = 42): # 1. Setup pretrained viT Weights weights = torchvision.models.ViT_B_16_Weights.DEFAULT # 2. Get transforms transforms = weights.transforms() # 3. Setup pretrained instance model = torchvision.models.vit_b_16(weights=weights) # 4. Freeze the base layers in the model (this will stop all layers from training) for params in model.parameters(): params.requires_grad = False # Set seeds for reproducibility torch.manual_seed(seeds) # 5. Modify the number of output layers model.heads = torch.nn.Sequential( torch.nn.Linear(in_features=768, out_features=num_classes, bias=True) ) return model, transforms