Update model.py

model.py

@@ -20,6 +20,7 @@ def create_resnet50_model(num_classes:int=10, # 4
     weights = torchvision.models.ResNet50_Weights.DEFAULT
     transforms = weights.transforms()
     model = torchvision.models.resnet50(weights=weights)
+    model.fc = nn.Linear(2048, 10)
 
     # 4. Freeze all layers in base model
     for param in model.parameters():