Spaces:

parneetsingh022
/

cat-vs-dog

Sleeping

parneetsingh022 commited on May 24, 2024

Commit

b41a5c0

verified ·

1 Parent(s): b0471d8

Update app.py

Files changed (1) hide show

app.py CHANGED Viewed

@@ -44,7 +44,8 @@ class CustomModel(nn.Module):
         return x
 model = CustomModel(input_shape=(3,128,128), num_classes=2)
-model.load_state_dict(torch.load('model.pth'))
 def greet(image):
     return "Hello " + "name" + "!!"

         return x
 model = CustomModel(input_shape=(3,128,128), num_classes=2)
+model.load_state_dict(torch.load('model.pth', map_location=torch.device('cpu')))
 def greet(image):
     return "Hello " + "name" + "!!"