Spaces:

bharatbk
/

custom_image_classification_1

Runtime error

bharatbk commited on Mar 2, 2023

Commit

e961972

1 Parent(s): 0b0621b

Update model.py

Files changed (1) hide show

model.py CHANGED Viewed

@@ -9,7 +9,7 @@ def create_effnetb2_model(num_classes:int=3, # default output classes =3 (pizza,
   # 1, 2, 3 create EffNetB2 pretrained weights, transforms and model
   weights=torchvision.models.EfficientNet_B2_Weights.DEFAULT
   transforms=weights.transforms()
-  model=torchvision.models.efficientnet_b2(weights=effnet_b2_weights)
   # 4. freeze all layers in the base model
   for params in model.parameters():

   # 1, 2, 3 create EffNetB2 pretrained weights, transforms and model
   weights=torchvision.models.EfficientNet_B2_Weights.DEFAULT
   transforms=weights.transforms()
+  model=torchvision.models.efficientnet_b2(weights=weights)
   # 4. freeze all layers in the base model
   for params in model.parameters():