Spaces:

RyanTietjen
/

Food_Classification_PyTorch

Build error

RyanTietjen commited on Aug 23, 2024

Commit

f03edb7

verified ·

1 Parent(s): 3f2e2db

Update app.py

Files changed (1) hide show

app.py CHANGED Viewed

@@ -18,7 +18,8 @@ model, transforms = vit_b_16(num_classes=101,
                              freeze_gradients=True,
                              unfreeze_blocks=0)
-model.load_state_dict(torch.load('vit_b_16_unfreeze_one_encoder_block_10_total_epochs.pth'))
 def predict_single_image(img):

                              freeze_gradients=True,
                              unfreeze_blocks=0)
+model.load_state_dict(torch.load('vit_b_16_unfreeze_one_encoder_block_10_total_epochs.pth',
+                                map_location=torch.device("cpu")))
 def predict_single_image(img):