Spaces:
Sleeping
Sleeping
xco2
commited on
Commit
·
6c10092
1
Parent(s):
fda67f2
改成cpu运行
Browse files
app.py
CHANGED
@@ -419,7 +419,7 @@ def merge_images(images: np.ndarray):
|
|
419 |
|
420 |
def get_models(device):
|
421 |
def modelLoad(model, model_path, data_parallel=False):
|
422 |
-
model.load_state_dict(torch.load(model_path), strict=True)
|
423 |
|
424 |
if data_parallel:
|
425 |
model = torch.nn.DataParallel(model)
|
|
|
419 |
|
420 |
def get_models(device):
|
421 |
def modelLoad(model, model_path, data_parallel=False):
|
422 |
+
model.load_state_dict(torch.load(model_path, map_location='cpu'), strict=True)
|
423 |
|
424 |
if data_parallel:
|
425 |
model = torch.nn.DataParallel(model)
|