if self.parameters.mode not in [base_layers.PREDICT, base_layers.TFLITE]:
      pooled_outputs = pooled_outputs * pooled_mask