Internal change.

PiperOrigin-RevId: 398264305

official/vision/beta/projects/volumetric_models/evaluation/segmentation_metrics.py

@@ -133,3 +133,4 @@ class DiceScore:
     if self._per_class_metric:
       for class_id in range(self._num_classes):
         self._dice_scores_per_class[class_id] = tf.Variable(0.0)
+        self._count_per_class[class_id] = tf.Variable(0.0)

official/vision/beta/projects/volumetric_models/tasks/semantic_segmentation_3d.py

@@ -342,4 +342,6 @@ class SemanticSegmentation3DTask(base_task.Task):
         metric_name = self.metrics[0].name + '/class_{0}'.format(
             i - 1) if i > 0 else self.metrics[0].name
         result.update({metric_name: metric_val})
+    else:
+      result.update({self.metrics[0].name: metric})
     return result