Merge pull request #2 from TreezzZ/master

Fix JMLH bugs
ymcidence · Apr 3, 2020 · 778dd1c · 778dd1c
2 parents 4be0237 + 2cccdc2
commit 778dd1c
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/model/jmlh.py b/model/jmlh.py
@@ -17,13 +17,13 @@ def call(self, inputs, training=None, mask=None):
         batch_size = tf.shape(inputs[1])[0]
         fc_1 = self.fc_1(inputs[1])
         eps = tf.ones([batch_size, self.bbn_dim]) / 2.
-        code = binary_activation.binary_activation(fc_1, eps)
+        code, _ = binary_activation.binary_activation(fc_1, eps)
         cls = self.fc_2(code)
         return code, tf.nn.sigmoid(fc_1), cls
 
     def run(self, feat_in):
         batch_size = tf.shape(feat_in)[0]
         fc_1 = self.fc_1(feat_in)
         eps = tf.ones([batch_size, self.bbn_dim]) / 2.
-        code = binary_activation.binary_activation(fc_1, eps)
+        code, _ = binary_activation.binary_activation(fc_1, eps)
         return code
diff --git a/train/jmlh_train.py b/train/jmlh_train.py
@@ -26,7 +26,7 @@ def train_step(model: JMLH, batch_data, opt: tf.optimizers.Optimizer):
         model_input = batch_data
         code, prob, cls_prob = model(model_input, training=True)
 
-        loss = jmlh_loss(prob, cls_prob, label=batch_data[3])
+        loss = jmlh_loss(prob, cls_prob, label=batch_data[2])
 
         gradient = tape.gradient(loss, sources=model.trainable_variables)
         opt.apply_gradients(zip(gradient, model.trainable_variables))