Manmay commited on
Commit
d99cbfa
1 Parent(s): a8317ed

make it work on cpu

Browse files
Files changed (1) hide show
  1. tortoise/models/hifigan_decoder.py +5 -2
tortoise/models/hifigan_decoder.py CHANGED
@@ -229,7 +229,10 @@ class HifiganGenerator(torch.nn.Module):
229
 
230
  if not conv_post_weight_norm:
231
  remove_weight_norm(self.conv_post)
232
-
 
 
 
233
  def forward(self, x, g=None):
234
  """
235
  Args:
@@ -287,7 +290,7 @@ class HifiganGenerator(torch.nn.Module):
287
  mode="linear",
288
  )
289
  g = g.unsqueeze(0)
290
- return self.forward(up_2.to("cuda"), g.transpose(1,2))
291
 
292
  def remove_weight_norm(self):
293
  print("Removing weight norm...")
 
229
 
230
  if not conv_post_weight_norm:
231
  remove_weight_norm(self.conv_post)
232
+
233
+ self.device = torch.device('cuda' if torch.cuda.is_available() else'cpu')
234
+ if torch.backends.mps.is_available():
235
+ self.device = torch.device('mps')
236
  def forward(self, x, g=None):
237
  """
238
  Args:
 
290
  mode="linear",
291
  )
292
  g = g.unsqueeze(0)
293
+ return self.forward(up_2.to(self.device), g.transpose(1,2))
294
 
295
  def remove_weight_norm(self):
296
  print("Removing weight norm...")