code: def test_mean_metric_broadcast(nan_strategy):
      """Check that weights gets broadcasted correctly when Nans are present."""
      metric = MeanMetric(nan_strategy=nan_strategy)
  
      x = torch.arange(5).float()
      x[1] = torch.tensor(float("nan"))
      w = torch.arange(5).float()
  
      metric.update(x, w)
      res = metric.compute()
      assert (
          round(res.item(), 4) == 3.2222
      )  # (0*0 + 2*2 + 3*3 + 4*4) / (0 + 2 + 3 + 4)
  
      x = torch.arange(5).float()
      w = torch.arange(5).float()
      w[1] = torch.tensor(float("nan"))
  
      metric.update(x, w)
      res = metric.compute()
      assert round(res.item(), 4) == 3.2222
  
imports: import numpy
  import pytest
  import torch
  from torchmetrics.aggregation import CatMetric, MaxMetric, MeanMetric, MinMetric, SumMetric
  from torchmetrics.collections import MetricCollection
  from unittests import BATCH_SIZE, NUM_BATCHES
  from unittests.helpers.testers import MetricTester
label_desc: Use image node_id 10 for calling a global function with example usage: test_mean_metric_broadcast(nan_strategy) without return types