masked_op#

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

Calculate operation on masked tensor.

Parameters:

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

  • op (str, optional) – 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