text: # Copyright The PyTorch Lightning team.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  import torch
  from typing import Any, Callable, Optional, Union
  
  from pytorch_lightning.metrics.metric import Metric
  
  
  class MeanSquaredError(Metric):
      """
      Computes mean squared error.
  
      Args:
          compute_on_step:
              Forward only calls ``update()`` and return None if this is set to False. default: True
          dist_sync_on_step:
              Synchronize metric state across processes at each ``forward()``
              before returning the value at the step. default: False
          process_group:
              Specify the process group on which synchronization is called. default: None (which selects the entire world)
  
      Example:
  
          >>> from pytorch_lightning.metrics import MeanSquaredError
          >>> target = torch.tensor([2.5, 5.0, 4.0, 8.0])
          >>> preds = torch.tensor([3.0, 5.0, 2.5, 7.0])
          >>> mean_squared_error = MeanSquaredError()
          >>> mean_squared_error(preds, target)
          tensor(0.8750)
  
      """
  
      def __init__(
          self,
          compute_on_step: bool = True,
          dist_sync_on_step: bool = False,
          process_group: Optional[Any] = None,
      ):
          super().__init__(
              compute_on_step=compute_on_step,
              dist_sync_on_step=dist_sync_on_step,
              process_group=process_group,
          )
  
          self.add_state("sum_squared_error", default=torch.tensor(0.0), dist_reduce_fx="sum")
          self.add_state("total", default=torch.tensor(0), dist_reduce_fx="sum")
  
      def update(self, preds: torch.Tensor, target: torch.Tensor):
          """
          Update state with predictions and targets.
  
          Args:
              preds: Predictions from model
              target: Ground truth values
          """
          self._check_same_shape(preds, target)
          squared_error = torch.pow(preds - target, 2)
  
          self.sum_squared_error += torch.sum(squared_error)
          self.total += target.numel()
  
      def compute(self):
          """
          Computes mean squared error over state.
          """
          return self.sum_squared_error / self.total
  