docstr: Computes MSE loss between preds and targets, subject to a mask.
  
  Args:
      preds: logits of shape [B1, ..., Bn, M]
      targets: index targets for preds of shape [B1, ..., Bn]
      masks: 0 means don't compute loss, 1 means compute loss
          shape [B1, ..., Bn]
  
  Returns:
      Scalar cross entropy loss.
docstr_tok: ["Computes","MSE","loss","between","preds","and","targets",",","subject","to","a","mask",".","Args",":","preds",":","logits","of","shape","[","B1",",","...",",","Bn",",","M","]","targets",":","index","targets","for","preds","of","shape","[","B1",",","...",",","Bn","]","masks",":","0","means","do","n't","compute","loss",",","1","means","compute","loss","shape","[","B1",",","...",",","Bn","]","Returns",":","Scalar","cross","entropy","loss","."]
code: def _masked_mse_loss(
      self,
      preds: TensorType,
      targets: TensorType,
      masks: TensorType,
  ) -> TensorType:
      """Computes MSE loss between preds and targets, subject to a mask.
  
      Args:
          preds: logits of shape [B1, ..., Bn, M]
          targets: index targets for preds of shape [B1, ..., Bn]
          masks: 0 means don't compute loss, 1 means compute loss
              shape [B1, ..., Bn]
  
      Returns:
          Scalar cross entropy loss.
      """
      losses = F.mse_loss(preds, targets, reduction="none")
      losses = losses * masks.reshape(
          *masks.shape, *([1] * (len(preds.shape) - len(masks.shape)))
      )
      return losses.mean()
  
imports: from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type, Union
  import gymnasium
  import numpy
  import tree
  from gymnasium.spaces import Box, Discrete
  from rllib_dt.dt.dt_torch_model import DTTorchModel
  from ray.rllib.models.catalog import ModelCatalog
  from ray.rllib.models.modelv2 import ModelV2
  from ray.rllib.models.torch.mingpt import configure_gpt_optimizer
  from ray.rllib.models.torch.torch_action_dist import TorchCategorical, TorchDeterministic, TorchDistributionWrapper
  from ray.rllib.policy.sample_batch import SampleBatch
  from ray.rllib.policy.torch_mixins import LearningRateSchedule
  from ray.rllib.policy.torch_policy_v2 import TorchPolicyV2
  from ray.rllib.utils.annotations import DeveloperAPI, PublicAPI, override
  from ray.rllib.utils.framework import try_import_torch
  from ray.rllib.utils.numpy import convert_to_numpy
  from ray.rllib.utils.threading import with_lock
  from ray.rllib.utils.torch_utils import apply_grad_clipping
  from ray.rllib.utils.typing import TensorShape, TensorStructType, TensorType, TrainerConfigDict
label_desc: Use image node_id 15 for calling the DTTorchPolicy obj's underlying member method code with example usage: obj._masked_mse_loss(preds, targets, masks) and returns: losses