fixed `sample_weight` shape bug in instance center losses

0fa40953 · srihari-humbarwadi · 5d7d327e · 0fa40953
Unverified Commit 0fa40953 authored Mar 20, 2022 by srihari-humbarwadi
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 2 deletions

official/vision/beta/projects/panoptic_maskrcnn/tasks/panoptic_deeplab.py ...beta/projects/panoptic_maskrcnn/tasks/panoptic_deeplab.py +2 -2

No files found.
--- a/official/vision/beta/projects/panoptic_maskrcnn/tasks/panoptic_deeplab.py
+++ b/official/vision/beta/projects/panoptic_maskrcnn/tasks/panoptic_deeplab.py
@@ -143,10 +143,10 @@ class PanopticDeeplabTask(base_task.Task):
        labels['semantic_weights'],
        dtype=model_outputs['instance_centers_heatmap'].dtype)
    things_mask = tf.cast(
-        labels['things_mask'],
+        tf.squeeze(labels['things_mask'], axis=3),
        dtype=model_outputs['instance_centers_heatmap'].dtype)
    valid_mask = tf.cast(
-        labels['valid_mask'],
+        tf.squeeze(labels['valid_mask'], axis=3),
        dtype=model_outputs['instance_centers_heatmap'].dtype)
    
    segmentation_loss = segmentation_loss_fn(