Spaces:

allispaul
/

foodvision_mini

Runtime error

allispaul commited on Jul 20, 2023

Commit

3119e39

•

1 Parent(s): e315722

fix typo

Files changed (1) hide show

model.py CHANGED Viewed

@@ -27,6 +27,6 @@ def create_effnetb2_model(num_classes: int = 3,
     torch.manual_seed(seed)
     model.classifier = nn.Sequential(
         nn.Dropout(p=0.3, inplace=True),
-        nn.Linear(in_features=1408, out_features=len(class_names))
     )
     return model, transforms

     torch.manual_seed(seed)
     model.classifier = nn.Sequential(
         nn.Dropout(p=0.3, inplace=True),
+        nn.Linear(in_features=1408, out_features=num_classes)
     )
     return model, transforms