Spaces:
Runtime error
Runtime error
| import torch | |
| import torchvision | |
| def create_vit(num_classes: int=6, | |
| 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=num_classes)) | |
| return model, transforms | |