Spaces:

adnankarim
/

waste_classification

Sleeping

adnankarim commited on Oct 25, 2024

Commit

31a9349

verified ·

1 Parent(s): 728b0e2

Update app.py

Files changed (1) hide show

app.py CHANGED Viewed

@@ -10,7 +10,7 @@ model = EfficientNet.from_name('efficientnet-b0')
 in_features = model._fc.in_features
 model._fc = torch.nn.Linear(in_features, 2)
-model.load_state_dict(torch.load('model_transfer.pt'))
 model.to(device)

 in_features = model._fc.in_features
 model._fc = torch.nn.Linear(in_features, 2)
+model.load_state_dict(torch.load('model_transfer.pt', map_location=torch.device('cpu')))
 model.to(device)