rewrite

model/base/module.py

@@ -30,7 +30,24 @@ def _activation(activation):
     elif activation == "Tanh":
         return nn.Tanh()
     else:
-        raise NotImplemented(activation)
+        raise NotImplementedError(f"{activation} not valid")
+
+
+class LinearBlock(nn.Module):
+    def __init__(self, in_features: int, out_features: int, bias=None, activation_type="ReLU", norm_type="NONE"):
+        super().__init__()
+
+        self.norm_type = norm_type
+        self.activation_type = activation_type
+
+        bias = _use_bias_checker(norm_type) if bias is None else bias
+        self.linear = nn.Linear(in_features, out_features, bias)
+
+        self.normalization = _normalization(norm_type, out_features)
+        self.activation = _activation(activation_type)
+
+    def forward(self, x):
+        return self.activation(self.normalization(self.linear(x)))
 
 
 class Conv2dBlock(nn.Module):