fixed mlp

ShomyLiu · Jul 21, 2021 · a119c1e · a119c1e
1 parent 29fa183
commit a119c1e
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/framework/prediction.py b/framework/prediction.py
@@ -155,4 +155,4 @@ def init_weight(self):
         nn.init.uniform_(self.fc.bias, a=0, b=0.2)
 
     def forward(self, feature, *args, **kwargs):
-        return F.relu(self.fc(feature)[0])
+        return F.relu(self.fc(feature))