Spaces:

fbrynpk
/

pneumonia-detection

Build error

fbrynpk commited on Oct 25, 2023

Commit

d3c48d3

1 Parent(s): 986860c

Fix model not set to device

Files changed (1) hide show

app.py CHANGED Viewed

@@ -24,7 +24,8 @@ EffNetB0_load_model, EffNetB0_transforms = create_effnet(
     model=torchvision.models.efficientnet_b0,
     in_features=1280,
     dropout=0.2,
-    out_features=2,
 )
 # Write a transform for image

     model=torchvision.models.efficientnet_b0,
     in_features=1280,
     dropout=0.2,
+    out_features=len(class_name),
+    device=device,
 )
 # Write a transform for image