import tensorflow as tf zero_out_module = tf.load_op_library('./zero_out.so') print(zero_out_module.zero_out([[1, 2], [3, 4]])) # Prints # array([[1, 0], [0, 0]], dtype=int32)