thunnai commited on
Commit
3be8ae4
·
1 Parent(s): 470e483

Add device method to BiCodecTokenizer

Browse files
Files changed (1) hide show
  1. sparktts/models/audio_tokenizer.py +5 -0
sparktts/models/audio_tokenizer.py CHANGED
@@ -53,6 +53,11 @@ class BiCodecTokenizer:
53
  f"{self.model_dir}/wav2vec2-large-xlsr-53"
54
  ).to(self.device)
55
  self.feature_extractor.config.output_hidden_states = True
 
 
 
 
 
56
 
57
  def get_ref_clip(self, wav: np.ndarray) -> np.ndarray:
58
  """Get reference audio clip for speaker embedding."""
 
53
  f"{self.model_dir}/wav2vec2-large-xlsr-53"
54
  ).to(self.device)
55
  self.feature_extractor.config.output_hidden_states = True
56
+
57
+ def to(self, device: torch.device):
58
+ self.device = device
59
+ self.model.to(self.device)
60
+ self.feature_extractor.to(self.device)
61
 
62
  def get_ref_clip(self, wav: np.ndarray) -> np.ndarray:
63
  """Get reference audio clip for speaker embedding."""