done with function learning tasks.

models/nac.py

@@ -28,13 +28,14 @@ class NAC(nn.Module):
 
         self.W_hat = Parameter(torch.Tensor(out_features, in_features))
         self.M_hat = Parameter(torch.Tensor(out_features, in_features))
-        self.W = F.tanh(self.W_hat) * F.sigmoid(self.M_hat)
+        self.W = Parameter(F.tanh(self.W_hat) * F.sigmoid(self.M_hat))
+        self.register_parameter('bias', None)
 
         init.kaiming_uniform_(self.W_hat, a=math.sqrt(5))
         init.kaiming_uniform_(self.M_hat, a=math.sqrt(5))
 
     def forward(self, input):
-        return F.linear(input, self.W, None)
+        return F.linear(input, self.W, self.bias)
 
     def extra_repr(self):
         return 'in_features={}, out_features={}'.format(self.in_features, self.out_features)

models/nalu.py

@@ -37,7 +37,7 @@ class NALU(nn.Module):
     def forward(self, input):
         a = self.nac(input)
         g = F.sigmoid(F.linear(input, self.G, None))
-        add_sub = a * g
+        add_sub = g * a
         log_input = torch.log(torch.abs(input) + self.eps)
         m = torch.exp(F.linear(log_input, self.W, None))
         mul_div = (1 - g) * m