Spaces:
Build error
Build error
init
Browse files
app.py
CHANGED
@@ -75,7 +75,7 @@ model_dir = os.path.join(os.path.abspath(os.path.dirname(__file__)), 'U-2-Net/sa
|
|
75 |
print("...load U2NET---173.6 MB")
|
76 |
net = U2NET(3,1)
|
77 |
|
78 |
-
net.load_state_dict(torch.load(model_dir))
|
79 |
# if torch.cuda.is_available():
|
80 |
# net.cuda()
|
81 |
net.eval()
|
|
|
75 |
print("...load U2NET---173.6 MB")
|
76 |
net = U2NET(3,1)
|
77 |
|
78 |
+
net.load_state_dict(torch.load(model_dir, map_location='cpu'))
|
79 |
# if torch.cuda.is_available():
|
80 |
# net.cuda()
|
81 |
net.eval()
|