Spaces:

Saugat20021
/

foodvision_mini

Sleeping

Saugat2002 commited on Jan 9

Commit

70f7338

1 Parent(s): 9baeefb

second try

Files changed (1) hide show

model.py CHANGED Viewed

@@ -7,7 +7,7 @@ def create_effnetb2_model(num_classes:int = 3, seed:int = 42):
   weights = torchvision.models.EfficientNet_B2_Weights.DEFAULT
   transforms = weights.transforms()
-  model = torchvision.models.efficientnet_b2(weights=weights).to(device)
   for param in model.parameters():
     param.requires_grad = False

   weights = torchvision.models.EfficientNet_B2_Weights.DEFAULT
   transforms = weights.transforms()
+  model = torchvision.models.efficientnet_b2(weights=weights)
   for param in model.parameters():
     param.requires_grad = False