Resolved comments 4 PR#44

facebookresearch · Nov 21, 2024 · 96f5d6f · 96f5d6f
1 parent 1baaf49
commit 96f5d6f
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 9 deletions.
diff --git a/sonar/models/mutox/builder.py b/sonar/models/mutox/builder.py
@@ -40,25 +40,29 @@ def __init__(
         self.config = config
         self.device, self.dtype = device, dtype
 
-    def build_model(self, activation=nn.ReLU()) -> MutoxClassifier:
+    def build_model(self) -> MutoxClassifier:
         model_h1 = nn.Sequential(
             nn.Dropout(0.01),
             nn.Linear(self.config.input_size, 512),
         )
 
         model_h2 = nn.Sequential(
-            activation,
+            nn.ReLU(),
             nn.Linear(512, 128),
         )
 
+        model_h3 = nn.Sequential(
+            nn.ReLU(),
+            nn.Linear(128, 1),
+        )
+
         model_all = nn.Sequential(
             model_h1,
             model_h2,
+            model_h3,
         )
 
-        return MutoxClassifier(
-            model_all,
-        ).to(
+        return MutoxClassifier(model_all,).to(
             device=self.device,
             dtype=self.dtype,
         )

diff --git a/sonar/models/mutox/classifier.py b/sonar/models/mutox/classifier.py
@@ -22,12 +22,13 @@ def __init__(
         self.model_all = model_all
 
     def forward(self, inputs: torch.Tensor, output_prob: bool = False) -> torch.Tensor:
+        outputs = self.model_all(inputs)
+
         if output_prob:
-            self.model_all.add_module("sigmoid", nn.Sigmoid())
-        else:
-            self.model_all.add_module("linear", nn.Linear(128, 1))
+            outputs = torch.sigmoid(outputs)
+
+        return outputs
 
-        return self.model_all(inputs)
 
 
 @dataclass