diff --git a/grad-cam.py b/grad-cam.py
index a175db7..593ed02 100644
--- a/grad-cam.py
+++ b/grad-cam.py
@@ -95,7 +95,7 @@ def grad_cam(input_model, image, category_index, layer_name):
                      output_shape = target_category_loss_output_shape))
 
     loss = K.sum(model.layers[-1].output)
-    conv_output =  [l for l in model.layers[0].layers if l.name is layer_name][0].output
+    conv_output =  [l for l in model.layers[0].layers if l.name == layer_name][0].output
     grads = normalize(K.gradients(loss, conv_output)[0])
     gradient_function = K.function([model.layers[0].input], [conv_output, grads])