Fix keras model compile for lifted_struct_loss

tensorflow_addons/losses/lifted.py

@@ -39,7 +39,6 @@ def lifted_struct_loss(labels, embeddings, margin=1.0):
     """
     # Reshape [batch_size] label tensor to a [batch_size, 1] label tensor.
     lshape = tf.shape(labels)
-    assert lshape.shape == 1
     labels = tf.reshape(labels, [lshape[0], 1])
 
     # Build pairwise squared distance matrix.

tensorflow_addons/losses/lifted_test.py

@@ -64,10 +64,14 @@ def testLiftedStruct(self):
         labels = np.random.randint(
             0, num_classes, size=(num_data)).astype(np.float32)
         # Reshape labels to compute adjacency matrix.
+        # pylint: disable=E1136
         labels_reshaped = np.reshape(labels, (labels.shape[0], 1))
+        # pylint: enable=E1136
 
         # Compute the loss in NP
+        # pylint: disable=E1111
         adjacency = np.equal(labels_reshaped, labels_reshaped.T)
+        # pylint: enable=E1111
         pdist_matrix = pairwise_distance_np(embedding)
         loss_np = 0.0
         num_constraints = 0.0
@@ -102,6 +106,13 @@ def testLiftedStruct(self):
         loss = cce_obj(y_true, y_pred)
         self.assertAlmostEqual(self.evaluate(loss), loss_np, 3)
 
+    def test_keras_model_compile(self):
+        model = tf.keras.models.Sequential([
+            tf.keras.layers.Input(shape=(784,)),
+            tf.keras.layers.Dense(10),
+        ])
+        model.compile(loss="Addons>lifted_struct_loss", optimizer="adam")
+
 
 if __name__ == '__main__':
     tf.test.main()