masked_op#

pytorch_forecasting.utils.masked_op(tensor: Tensor, op: str = 'mean', dim: int = 0, mask: Optional[Tensor] = None) Tensor[source]#

Calculate operation on masked tensor.

Parameters
  • tensor (torch.Tensor) – tensor to conduct operation over

  • op (str) – operation to apply. One of [“mean”, “sum”]. Defaults to “mean”.

  • dim (int, optional) – dimension to average over. Defaults to 0.

  • mask (torch.Tensor, optional) – boolean mask to apply (True=will take mean, False=ignore). Masks nan values by default.

Returns

tensor with averaged out dimension

Return type

torch.Tensor