swimmiing commited on
Commit
833fa47
1 Parent(s): 8fee4bf

Fix cpu only v2

Browse files
Files changed (2) hide show
  1. app.py +1 -1
  2. modules/models.py +1 -1
app.py CHANGED
@@ -9,7 +9,7 @@ from PIL import Image
9
 
10
 
11
  def greet(image, audio):
12
- device = "cpu"
13
 
14
  # Get model
15
  model_conf_file = f'./config/model/ACL_ViT16.yaml'
 
9
 
10
 
11
  def greet(image, audio):
12
+ device = torch.device('cpu')
13
 
14
  # Get model
15
  model_conf_file = f'./config/model/ACL_ViT16.yaml'
modules/models.py CHANGED
@@ -289,6 +289,6 @@ class ACL(nn.Module):
289
  Args:
290
  model_dir (str): Directory to load the model from.
291
  """
292
- ckp = torch.load(model_dir)
293
  self.audio_proj.load_state_dict(ckp['audio_proj'])
294
  self.masker_i.load_state_dict(ckp['masker_i'])
 
289
  Args:
290
  model_dir (str): Directory to load the model from.
291
  """
292
+ ckp = torch.load(model_dir, map_location=self.device)
293
  self.audio_proj.load_state_dict(ckp['audio_proj'])
294
  self.masker_i.load_state_dict(ckp['masker_i'])