fix a type, the input for center loss should be embedding layer

2807604035 · Oct 29, 2016 · c03bca7 · c03bca7
1 parent e26ef80
commit c03bca7
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/train.py b/train.py
@@ -71,7 +71,7 @@ def get_symbol(batchsize=64):
 
 
     ce_loss = mx.symbol.SoftmaxOutput(data=fc2, label=softmax_label, name='softmax')
-    center_loss_ = mx.symbol.Custom(data=fc2, label=center_label, name='center_loss_', op_type='centerloss',\
+    center_loss_ = mx.symbol.Custom(data=embedding, label=center_label, name='center_loss_', op_type='centerloss',\
             num_class=10, alpha=0.5, scale=0.01, batchsize=batchsize)
     center_loss = mx.symbol.MakeLoss(name='center_loss', data=center_loss_)
     mlp = mx.symbol.Group([ce_loss, center_loss])