just use device off the raw audio tensor passed in

lucidrains · Nov 22, 2023 · 93a07dd · 93a07dd
1 parent 7aebf44
commit 93a07dd
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 14 deletions.
diff --git a/bs_roformer/bs_roformer.py b/bs_roformer/bs_roformer.py
@@ -349,10 +349,6 @@ def __init__(
             normalized = multi_stft_normalized
         )
 
-    @property
-    def device(self):
-        return next(self.parameters()).device
-
     def forward(
         self,
         raw_audio,
@@ -371,6 +367,8 @@ def forward(
         d - feature dimension
         """
 
+        device = raw_audio.device
+
         if raw_audio.ndim == 2:
             raw_audio = rearrange(raw_audio, 'b t -> b 1 t')
 
@@ -381,7 +379,7 @@ def forward(
 
         raw_audio, batch_audio_channel_packed_shape = pack_one(raw_audio, '* t')
 
-        stft_window = self.stft_window_fn(device = self.device)
+        stft_window = self.stft_window_fn(device = device)
 
         stft_repr = torch.stft(raw_audio, **self.stft_kwargs, window = stft_window, return_complex = True)
         stft_repr = torch.view_as_real(stft_repr)
@@ -462,7 +460,7 @@ def forward(
                 n_fft = max(window_size, self.multi_stft_n_fft),  # not sure what n_fft is across multi resolution stft
                 win_length = window_size,
                 return_complex = True,
-                window = self.multi_stft_window_fn(window_size, device = self.device),
+                window = self.multi_stft_window_fn(window_size, device = device),
                 **self.multi_stft_kwargs,
             )
 

diff --git a/bs_roformer/mel_band_roformer.py b/bs_roformer/mel_band_roformer.py
@@ -383,10 +383,6 @@ def __init__(
 
         self.match_input_audio_length = match_input_audio_length
 
-    @property
-    def device(self):
-        return next(self.parameters()).device
-
     def forward(
         self,
         raw_audio,
@@ -405,6 +401,8 @@ def forward(
         d - feature dimension
         """
 
+        device = raw_audio.device
+
         if raw_audio.ndim == 2:
             raw_audio = rearrange(raw_audio, 'b t -> b 1 t')
 
@@ -418,7 +416,7 @@ def forward(
 
         raw_audio, batch_audio_channel_packed_shape = pack_one(raw_audio, '* t')
 
-        stft_window = self.stft_window_fn(device = self.device)
+        stft_window = self.stft_window_fn(device = device)
 
         stft_repr = torch.stft(raw_audio, **self.stft_kwargs, window = stft_window, return_complex = True)
         stft_repr = torch.view_as_real(stft_repr)
@@ -428,7 +426,7 @@ def forward(
 
         # index out all frequencies for all frequency ranges across bands ascending in one go
 
-        batch_arange = torch.arange(batch, device = self.device)[..., None]
+        batch_arange = torch.arange(batch, device = device)[..., None]
 
         # account for stereo
 
@@ -522,7 +520,7 @@ def forward(
                 n_fft = max(window_size, self.multi_stft_n_fft),  # not sure what n_fft is across multi resolution stft
                 win_length = window_size,
                 return_complex = True,
-                window = self.multi_stft_window_fn(window_size, device = self.device),
+                window = self.multi_stft_window_fn(window_size, device = device),
                 **self.multi_stft_kwargs,
             )
 

diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'BS-RoFormer',
   packages = find_packages(exclude=[]),
-  version = '0.3.8',
+  version = '0.3.9',
   license='MIT',
   description = 'BS-RoFormer - Band-Split Rotary Transformer for SOTA Music Source Separation',
   author = 'Phil Wang',