diff --git a/tensorflow_addons/losses/lifted.py b/tensorflow_addons/losses/lifted.py
index de2fb07de3..ef69f5cab7 100644
--- a/tensorflow_addons/losses/lifted.py
+++ b/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.
diff --git a/tensorflow_addons/losses/lifted_test.py b/tensorflow_addons/losses/lifted_test.py
index 0f2b9fda88..2191639eea 100644
--- a/tensorflow_addons/losses/lifted_test.py
+++ b/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()