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) – 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