The problem most likely is to do with how you handle the scalar values in that computation. You have to be careful not to use TF functions for computing the scalar values, because they have very different type coercion rules. Either use python or numpy for the scalar parts and only use the TF functions for the tensors. Here’s a thread which discusses issues related to this.
2 Likes