device bug

Author: RichieHakim

face_rhythm/pipelines.py

@@ -222,7 +222,7 @@ def pipeline_basic(params):
         ## - `downsample_factor`: How much to downsample the spectrogram by in time.
         ## - `return_complex`: Whether or not to return the complex spectrogram. Generally set to False unless you want to try something fancy.
 
-        params['VQT_Analyzer']['params_VQT']['DEVICE_compute'] = fr.helpers.set_device(use_GPU=use_GPU)
+        params['VQT_Analyzer']['device'] = fr.helpers.set_device(use_GPU=use_GPU)
 
         spec = fr.spectral_analysis.VQT_Analyzer(**params['VQT_Analyzer'])
 

face_rhythm/spectral_analysis.py

@@ -51,7 +51,7 @@ def __init__(
             'plot_pref': False,
         },
         batch_size: int=10,
-        DEVICE_compute='cpu',
+        device='cpu',
         normalization_factor: float=0.99,
         spectrogram_exponent: float=1.0,
         one_over_f_exponent: float=1.0,
@@ -62,7 +62,7 @@ def __init__(
         ## Set attributes
         self._params_VQT = params_VQT
         self._batch_size = int(batch_size)
-        self._DEVICE_compute = DEVICE_compute
+        self._device = device
         self._normalization_factor = float(normalization_factor)
         self._spectrogram_exponent = float(spectrogram_exponent)
         self._one_over_f_exponent = float(one_over_f_exponent)
@@ -84,7 +84,7 @@ def __init__(
         self.config = {
             'params_VQT': params_VQT,
             'batch_size': batch_size,
-            'DEVICE_compute': DEVICE_compute,
+            'device': device,
             'normalization_factor': normalization_factor,
             'spectrogram_exponent': spectrogram_exponent,
             'one_over_f_exponent': one_over_f_exponent,
@@ -135,9 +135,9 @@ def transform(self, points_tracked: np.ndarray, point_positions: np.ndarray):
         freqs = self.vqt_model.freqs
         xAxis = self.vqt_model.get_xAxis(points_tracked.shape[-1])
         ### send vqt_model to device
-        self.vqt_model.to(self._DEVICE_compute)
+        self.vqt_model.to(self._device)
         spec = torch.cat([
-            self.vqt_model(p.to(self._DEVICE_compute)).cpu()
+            self.vqt_model(p.to(self._device)).cpu()
             for p in tqdm(helpers.make_batches(points_tracked, batch_size=self._batch_size), disable=not self._verbose > 1, desc='Computing spectrograms', leave=True, position=0, total=int(math.ceil(points_tracked.shape[0] / self._batch_size)))
         ], dim=0)
         self.vqt_model.to('cpu')