import torch import torchvision def create_vit(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 = torch.nn.Sequential(torch.nn.LayerNorm(normalized_shape=768), torch.nn.Linear(in_features=768, out_features=1)) return model, transforms