Switch S4D to ml4gw import

libs/architectures/architectures/networks/__init__.py

@@ -1,3 +1,2 @@
-from .s4 import S4Model
 from .wavenet import WaveNet
 from .xylophone import Xylophone

libs/architectures/architectures/supervised.py

@@ -1,10 +1,11 @@
 from typing import Literal, Optional
 
 from architectures import Architecture
-from architectures.networks import S4Model, WaveNet, Xylophone
+from architectures.networks import WaveNet, Xylophone
 from jaxtyping import Float
 from ml4gw.nn.resnet.resnet_1d import NormLayer, ResNet1D
 from ml4gw.nn.resnet.resnet_2d import ResNet2D
+from ml4gw.nn.ssm.s4d import S4Model
 from torch import Tensor
 import torch
 

libs/architectures/pyproject.toml

@@ -7,7 +7,7 @@ requires-python = ">=3.10,<3.13"
 license = "MIT"
 dependencies = [
     "einops>=0.8,<0.9",
-    "ml4gw>=0.7.2",
+    "ml4gw>=0.8.3",
     "h5py>=3.9.0,<4",
     "numpy~=1.26",
 ]