Check for tf.Tensor type in tpu_util.write_to_variable.

PiperOrigin-RevId: 321891941

Author: shraman-rc

morph_net/framework/tpu_util.py

@@ -97,6 +97,9 @@ def maybe_convert_to_variable(tensor):
 
 def write_to_variable(tensor, fail_if_exists=True):
   """Saves a tensor for later retrieval on CPU."""
+  if not isinstance(tensor, tf.Tensor):
+    raise ValueError('Expected tf.Tensor but got {}'.format(type(tensor)))
+
   # Only relevant for debugging.
   debug_name = 'tpu_util__' + tensor.name.split(':')[0]
 

morph_net/framework/tpu_util_test.py

@@ -104,5 +104,11 @@ def test_write_to_variable(self):
       tpu_util.write_to_variable(bar)
       self.assertLen(set(tpu_util.var_store.values()), 2)
 
+  def test_write_to_variable_check_inputs(self):
+    variable = tf.get_variable('x', shape=[1], dtype=tf.float32)
+    with self.assertRaises(ValueError):
+      tpu_util.write_to_variable(variable)
+
+
 if __name__ == '__main__':
   tf.test.main()