zero_out.py 178 B

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