[FIX] Numeric type error for tf.zeros

wbenbihi · Aug 20, 2022 · 56d9ae9 · 56d9ae9
1 parent 68ab58c
commit 56d9ae9
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/hourglass_tensorflow/datasets/transformation.py b/hourglass_tensorflow/datasets/transformation.py
@@ -192,7 +192,7 @@ def tf_train_map_heatmaps(
                 joint[:2], stddev_tensor, shape_tensor, precision=precision
             )
             if joint[2] == 1.0
-            else tf.zeros(shape_tensor, dtype=precision)
+            else tf.zeros(tf.cast(shape_tensor, dtype=tf.dtypes.int32), dtype=precision)
         ),
         elems=joints,
         dtype=precision,