xco2 commited on
Commit
6c10092
·
1 Parent(s): fda67f2

改成cpu运行

Browse files
Files changed (1) hide show
  1. app.py +1 -1
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)