Update model

JusperLee · Sep 28, 2024 · 0b88cc3 · 0b88cc3
1 parent e2d6546
commit 0b88cc3
Showing 1 changed file with 6 additions and 5 deletions.
diff --git a/look2hear/models/SPMamba.py b/look2hear/models/SPMamba.py
@@ -13,10 +13,11 @@
 from torch_complex.tensor import ComplexTensor
 is_torch_1_9_plus = V(torch.__version__) >= V("1.9.0")
 
-from stft_tfgn import Stft
-from complex_utils import is_torch_complex_tensor
-from complex_utils import new_complex_like
-from get_layer_from_string import get_layer
+from ..layers import Stft
+from ..utils.complex_utils import is_torch_complex_tensor
+from ..utils.complex_utils import new_complex_like
+from ..utils.get_layer_from_string import get_layer
+from .base_model import BaseModel
 
 from functools import partial
 from mamba_ssm.modules.mamba_simple import Mamba, Block
@@ -400,7 +401,7 @@ def streaming_merge(self, chunks, ilens=None):
 
 
 
-class SPMamba(nn.Module):
+class SPMamba(BaseModel):
     def __init__(
         self,
         input_dim,