fix Cast bug

Alter-Y · Aug 5, 2022 · 46a0aaf · 46a0aaf
1 parent 56cea43
commit 46a0aaf
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/layers/common_layers.py b/layers/common_layers.py
@@ -285,8 +285,8 @@ def __call__(self, inputs):
                     inputs[i] = tf.cast(input[i], dtype=self.tf_cast_map[self.cast_to])
         else:
             if isinstance(inputs, np.ndarray) or isinstance(inputs, np.generic):
-                inputs[i] = np_cast_op(inputs)
+                inputs = np_cast_op(inputs)
             else:
-                inputs[i] = tf.cast(inputs, dtype=self.tf_cast_map[self.cast_to])
+                inputs = tf.cast(inputs, dtype=self.tf_cast_map[self.cast_to])
 
         return inputs