Error for 2d weights custom Keras metric in TFMA

gistfile1.txt

raceback (most recent call last):
  File "/home/user/project/project/main.py", line 269, in <module>
    LocalDagRunner().run(pipeline)
  File "/home/user/.local/share/virtualenvs/project-xv7hiZ8W/lib/python3.8/site-packages/tfx/orchestration/local/local_dag_runner.py", line 129, in run
    component_launcher.launch()
  File "/home/user/.local/share/virtualenvs/project-xv7hiZ8W/lib/python3.8/site-packages/tfx/orchestration/portable/launcher.py", line 412, in launch
    executor_output = self._run_executor(execution_info)
  File "/home/user/.local/share/virtualenvs/project-xv7hiZ8W/lib/python3.8/site-packages/tfx/orchestration/portable/launcher.py", line 306, in _run_executor
    executor_output = self._executor_operator.run_executor(execution_info)
  File "/home/user/.local/share/virtualenvs/project-xv7hiZ8W/lib/python3.8/site-packages/tfx/orchestration/portable/python_executor_operator.py", line 111, in run_executor
    result = executor.Do(execution_info.input_dict, output_dict,
  File "/home/user/project/project/project/custom_components/evaluator/my_executor.py", line 253, in Do
    (
  File "/home/user/.local/share/virtualenvs/project-xv7hiZ8W/lib/python3.8/site-packages/apache_beam/pipeline.py", line 580, in __exit__
    self.result = self.run()
  File "/home/user/.local/share/virtualenvs/project-xv7hiZ8W/lib/python3.8/site-packages/apache_beam/pipeline.py", line 559, in run
    return self.runner.run_pipeline(self, self._options)
  File "/home/user/.local/share/virtualenvs/project-xv7hiZ8W/lib/python3.8/site-packages/apache_beam/runners/direct/direct_runner.py", line 133, in run_pipeline
    return runner.run_pipeline(pipeline, options)
  File "/home/user/.local/share/virtualenvs/project-xv7hiZ8W/lib/python3.8/site-packages/apache_beam/runners/portability/fn_api_runner/fn_runner.py", line 182, in run_pipeline
    self._latest_run_result = self.run_via_runner_api(
  File "/home/user/.local/share/virtualenvs/project-xv7hiZ8W/lib/python3.8/site-packages/apache_beam/runners/portability/fn_api_runner/fn_runner.py", line 193, in run_via_runner_api
    return self.run_stages(stage_context, stages)
  File "/home/user/.local/share/virtualenvs/project-xv7hiZ8W/lib/python3.8/site-packages/apache_beam/runners/portability/fn_api_runner/fn_runner.py", line 357, in run_stages
    stage_results = self._run_stage(
  File "/home/user/.local/share/virtualenvs/project-xv7hiZ8W/lib/python3.8/site-packages/apache_beam/runners/portability/fn_api_runner/fn_runner.py", line 548, in _run_stage
    last_result, deferred_inputs, fired_timers = self._run_bundle(
  File "/home/user/.local/share/virtualenvs/project-xv7hiZ8W/lib/python3.8/site-packages/apache_beam/runners/portability/fn_api_runner/fn_runner.py", line 594, in _run_bundle
    result, splits = bundle_manager.process_bundle(
  File "/home/user/.local/share/virtualenvs/project-xv7hiZ8W/lib/python3.8/site-packages/apache_beam/runners/portability/fn_api_runner/fn_runner.py", line 896, in process_bundle
    result_future = self._worker_handler.control_conn.push(process_bundle_req)
  File "/home/user/.local/share/virtualenvs/project-xv7hiZ8W/lib/python3.8/site-packages/apache_beam/runners/portability/fn_api_runner/worker_handlers.py", line 380, in push
    response = self.worker.do_instruction(request)
  File "/home/user/.local/share/virtualenvs/project-xv7hiZ8W/lib/python3.8/site-packages/apache_beam/runners/worker/sdk_worker.py", line 606, in do_instruction
    return getattr(self, request_type)(
  File "/home/user/.local/share/virtualenvs/project-xv7hiZ8W/lib/python3.8/site-packages/apache_beam/runners/worker/sdk_worker.py", line 644, in process_bundle
    bundle_processor.process_bundle(instruction_id))
  File "/home/user/.local/share/virtualenvs/project-xv7hiZ8W/lib/python3.8/site-packages/apache_beam/runners/worker/bundle_processor.py", line 1005, in process_bundle
    op.finish()
  File "apache_beam/runners/worker/operations.py", line 1080, in apache_beam.runners.worker.operations.PGBKCVOperation.finish
  File "apache_beam/runners/worker/operations.py", line 1083, in apache_beam.runners.worker.operations.PGBKCVOperation.finish
  File "apache_beam/runners/worker/operations.py", line 1098, in apache_beam.runners.worker.operations.PGBKCVOperation.output_key
  File "/home/user/.local/share/virtualenvs/project-xv7hiZ8W/lib/python3.8/site-packages/tensorflow_model_analysis/evaluators/metrics_plots_and_validations_evaluator.py", line 396, in compact
    return super(_ComputationsCombineFn, self).compact(accumulator)
  File "/home/user/.local/share/virtualenvs/project-xv7hiZ8W/lib/python3.8/site-packages/apache_beam/transforms/combiners.py", line 756, in compact
    return [
  File "/home/user/.local/share/virtualenvs/project-xv7hiZ8W/lib/python3.8/site-packages/apache_beam/transforms/combiners.py", line 757, in <listcomp>
    c.compact(a, *args, **kwargs) for c,
  File "/home/user/.local/share/virtualenvs/project-xv7hiZ8W/lib/python3.8/site-packages/tensorflow_model_analysis/metrics/tf_metric_wrapper.py", line 568, in compact
    self._process_batch(accumulator)
  File "/home/user/.local/share/virtualenvs/project-xv7hiZ8W/lib/python3.8/site-packages/tensorflow_model_analysis/metrics/tf_metric_wrapper.py", line 497, in _process_batch
    accumulator.add_weights(output_index, metric_index,
  File "/home/user/.local/share/virtualenvs/project-xv7hiZ8W/lib/python3.8/site-packages/tensorflow_model_analysis/metrics/tf_metric_accumulators.py", line 120, in add_weights
    self._weights[output_index][metric_index] = np.add(cur_weights, weights)
ValueError: operands could not be broadcast together with shapes (1,937,30208) (1,73,30208)